1558 lines
80 KiB
Python
1558 lines
80 KiB
Python
|
"""RAG 管道：素材导入、增量重建与检索。"""
|
|||
|
|
import shutil
|
|||
|
|
import hashlib
|
|||
|
|
import uuid
|
|||
|
|
import logging
|
|||
|
|
import threading
|
|||
|
|
import time
|
|||
|
|
import os
|
|||
|
|
import re
|
|||
|
|
import unicodedata
|
|||
|
|
from collections import deque
|
|||
|
|
from concurrent.futures import ThreadPoolExecutor
|
|||
|
|
from pathlib import Path
|
|||
|
|
from sqlmodel import Session, select
|
|||
|
|
from sqlalchemy import text, or_
|
|||
|
|
|
|||
|
|
from ..config import KB_DIR
|
|||
|
|
from ..models import Campaign, KnowledgeBase, KnowledgeUsage, KnowledgeCategory, KnowledgeFile, now
|
|||
|
|
from ..billing import estimate_tokens, record_usage_event
|
|||
|
|
from . import backend, parse, chunk, embed, store, dedup
|
|||
|
|
from . import hybrid
|
|||
|
|
from .quarantine import file_exclusions
|
|||
|
|
from .context import source_excerpt
|
|||
|
|
from .. import perf_metrics
|
|||
|
|
|
|||
|
|
logger = logging.getLogger("yawen.rag.pipeline")
|
|||
|
|
|
|||
|
|
# Reuse a bounded set of recall workers across requests. Creating a fresh
|
|||
|
|
# executor for every multi-library query defeated thread-local SQLite/FTS
|
|||
|
|
# reuse and could create an unbounded number of threads under concurrent load.
|
|||
|
|
# The queue is intentionally bounded by the executor's fixed worker count;
|
|||
|
|
# callers still get the same per-request fan-out semantics.
|
|||
|
|
_recall_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="rag-recall")
|
|||
|
|
_PREWARM_LOCK = threading.Lock()
|
|||
|
|
_PREWARMING: set[int] = set()
|
|||
|
|
|
|||
|
|
# A bounded, process-local rolling window keeps retrieval latency observable
|
|||
|
|
# without adding a database write to every query. The worker process logs a
|
|||
|
|
# compact P50/P95 summary periodically; callers still receive per-query phase
|
|||
|
|
# timings in ``retrieval_timings`` below.
|
|||
|
|
_RETRIEVAL_METRICS_LOCK = threading.Lock()
|
|||
|
|
_RETRIEVAL_METRICS: deque[dict[str, float | int | bool]] = deque(maxlen=512)
|
|||
|
|
_RETRIEVAL_METRICS_LOG_EVERY = 25
|
|||
|
|
_LAST_RETRIEVAL_METRIC: dict | None = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _record_query_embedding_usage(*, idempotency_key: str, user_id: int,
|
|||
|
|
organization_id: int, campaign_id: int,
|
|||
|
|
trace_id: str, model: str, quantity: float,
|
|||
|
|
metadata: dict) -> None:
|
|||
|
|
"""Write embedding telemetry in its own short transaction.
|
|||
|
|
|
|||
|
|
Retrieval is read-heavy and can fan out across several content tasks. A
|
|||
|
|
usage event flushed through the caller's long-lived transaction made those
|
|||
|
|
readers compete for SQLite's single writer and surfaced ``database is
|
|||
|
|
locked`` as a false retrieval failure. The event is accounting metadata;
|
|||
|
|
it must never make a successful retrieval fail.
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
from ..db import engine
|
|||
|
|
with Session(engine) as usage_session:
|
|||
|
|
# Telemetry must never wait for the full 30 s application timeout
|
|||
|
|
# behind a content-task write transaction. If SQLite is busy,
|
|||
|
|
# skip this non-critical event quickly and keep retrieval usable.
|
|||
|
|
usage_session.connection().exec_driver_sql("PRAGMA busy_timeout=500")
|
|||
|
|
record_usage_event(
|
|||
|
|
usage_session, idempotency_key=idempotency_key,
|
|||
|
|
user_id=user_id, organization_id=organization_id,
|
|||
|
|
campaign_id=campaign_id, trace_id=trace_id, model=model,
|
|||
|
|
service_type="embedding", unit="token", quantity=quantity,
|
|||
|
|
status="estimated", metadata=metadata)
|
|||
|
|
usage_session.commit()
|
|||
|
|
except Exception as exc: # pragma: no cover - telemetry must be best effort
|
|||
|
|
logger.warning("Embedding 用量记录失败（不阻断知识库检索）：%s", exc)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def retrieval_metrics_snapshot() -> dict:
|
|||
|
|
"""Return process-local retrieval latency percentiles for diagnostics.
|
|||
|
|
|
|||
|
|
This is intentionally read-only and bounded. It is safe to expose from a
|
|||
|
|
diagnostics endpoint later, while keeping the search API unchanged.
|
|||
|
|
"""
|
|||
|
|
with _RETRIEVAL_METRICS_LOCK:
|
|||
|
|
rows = list(_RETRIEVAL_METRICS)
|
|||
|
|
if not rows:
|
|||
|
|
return {"samples": 0, "p50_ms": 0.0, "p95_ms": 0.0,
|
|||
|
|
"cache_hit_rate": 0.0, "phases": {}}
|
|||
|
|
|
|||
|
|
def percentile(values: list[float], fraction: float) -> float:
|
|||
|
|
values = sorted(values)
|
|||
|
|
if not values:
|
|||
|
|
return 0.0
|
|||
|
|
index = min(len(values) - 1, max(0, int(round((len(values) - 1) * fraction))))
|
|||
|
|
return round(values[index], 2)
|
|||
|
|
|
|||
|
|
total = [float(row.get("total_ms", 0.0)) for row in rows]
|
|||
|
|
phase_names = ("validation_ms", "embedding_ms", "recall_ms", "postprocess_ms")
|
|||
|
|
phases = {name: {"p50_ms": percentile([float(row.get(name, 0.0)) for row in rows], 0.5),
|
|||
|
|
"p95_ms": percentile([float(row.get(name, 0.0)) for row in rows], 0.95)}
|
|||
|
|
for name in phase_names}
|
|||
|
|
return {"samples": len(rows), "p50_ms": percentile(total, 0.5),
|
|||
|
|
"p95_ms": percentile(total, 0.95),
|
|||
|
|
"cache_hit_rate": round(sum(bool(row.get("cache_hit")) for row in rows) / len(rows), 4),
|
|||
|
|
"phases": phases}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _record_retrieval_metrics(timings: dict, result_count: int) -> None:
|
|||
|
|
"""Record a bounded sample and emit a periodic percentile summary."""
|
|||
|
|
global _RETRIEVAL_METRICS_LOG_EVERY, _LAST_RETRIEVAL_METRIC
|
|||
|
|
_LAST_RETRIEVAL_METRIC = {**timings, "result_count": int(result_count)}
|
|||
|
|
perf_metrics.ingest_rag(_LAST_RETRIEVAL_METRIC)
|
|||
|
|
with _RETRIEVAL_METRICS_LOCK:
|
|||
|
|
_RETRIEVAL_METRICS.append({**timings, "result_count": int(result_count)})
|
|||
|
|
count = len(_RETRIEVAL_METRICS)
|
|||
|
|
should_log = count == 1 or count % _RETRIEVAL_METRICS_LOG_EVERY == 0
|
|||
|
|
if should_log:
|
|||
|
|
summary = retrieval_metrics_snapshot()
|
|||
|
|
# ``rag_worker`` deliberately has no logging configuration because its
|
|||
|
|
# stdout is a JSON protocol. WARNING is handled by Python's fallback
|
|||
|
|
# logger, so the summary remains visible in rag-worker.log without
|
|||
|
|
# touching that protocol or requiring a global logging setup.
|
|||
|
|
logger.warning("RAG检索指标：样本=%s，P50=%.2fms，P95=%.2fms，缓存命中率=%.1f%%，阶段=%s",
|
|||
|
|
summary["samples"], summary["p50_ms"], summary["p95_ms"],
|
|||
|
|
summary["cache_hit_rate"] * 100, summary["phases"])
|
|||
|
|
|
|||
|
|
def last_retrieval_metric() -> dict:
|
|||
|
|
return dict(_LAST_RETRIEVAL_METRIC or {})
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _schedule_lexical_prewarm(kb_id: int, organization_id: int = 1) -> None:
|
|||
|
|
"""Warm the default full-library FTS scope after a successful import.
|
|||
|
|
|
|||
|
|
FTS construction can scan a large library, so it must never delay the
|
|||
|
|
upload response or the first user query. The hybrid module's file lock
|
|||
|
|
makes retries and multiple API workers safe.
|
|||
|
|
"""
|
|||
|
|
kid = int(kb_id)
|
|||
|
|
with _PREWARM_LOCK:
|
|||
|
|
if kid in _PREWARMING:
|
|||
|
|
return
|
|||
|
|
_PREWARMING.add(kid)
|
|||
|
|
|
|||
|
|
def run() -> None:
|
|||
|
|
started = __import__("time").perf_counter()
|
|||
|
|
try:
|
|||
|
|
from . import hybrid
|
|||
|
|
hybrid.prewarm_index(int(organization_id or 1), kid)
|
|||
|
|
logger.info("知识库全文索引已后台预热：知识库 %s，%.3f 秒", kid,
|
|||
|
|
__import__("time").perf_counter() - started)
|
|||
|
|
except Exception as exc:
|
|||
|
|
# Vector import has already succeeded; a failed warmup only means
|
|||
|
|
# the next query will retry it under the normal index lock.
|
|||
|
|
logger.warning("知识库全文索引预热失败：知识库 %s，原因：%s", kid, exc)
|
|||
|
|
finally:
|
|||
|
|
with _PREWARM_LOCK:
|
|||
|
|
_PREWARMING.discard(kid)
|
|||
|
|
|
|||
|
|
threading.Thread(target=run, name=f"rag-prewarm-{kid}", daemon=True).start()
|
|||
|
|
|
|||
|
|
SUPPORTED_SUFFIXES = {".doc", ".docx", ".srt", ".md", ".txt", ".json", ".jsonl", ".pdf"}
|
|||
|
|
# 音视频后缀（whisper 本地转写）
|
|||
|
|
MEDIA_SUFFIXES = {".mp4", ".mov", ".m4a", ".mp3", ".wav", ".flv", ".mkv", ".avi", ".webm", ".aac", ".ogg", ".flac"}
|
|||
|
|
|
|||
|
|
def _source_dir(kb: KnowledgeBase) -> Path:
|
|||
|
|
return Path(kb.source_dir) if kb.source_dir else KB_DIR / f"{kb.id}_{kb.name}"
|
|||
|
|
|
|||
|
|
def _disk_files(source: Path) -> list[Path]:
|
|||
|
|
if not source.exists() or not source.is_dir():
|
|||
|
|
return []
|
|||
|
|
return sorted(path for path in source.iterdir()
|
|||
|
|
if path.is_file() and path.suffix.lower() in SUPPORTED_SUFFIXES | MEDIA_SUFFIXES)
|
|||
|
|
|
|||
|
|
def reconcile_kb_files(session: Session, kb_id: int | None = None) -> dict:
|
|||
|
|
"""启动时对账物理文件、文件表和向量计数，不自动重建大库。"""
|
|||
|
|
query = select(KnowledgeBase)
|
|||
|
|
if kb_id is not None:
|
|||
|
|
query = query.where(KnowledgeBase.id == kb_id)
|
|||
|
|
repaired = []
|
|||
|
|
for kb in session.exec(query).all():
|
|||
|
|
source = _source_dir(kb)
|
|||
|
|
files = _disk_files(source)
|
|||
|
|
known = {item.file_name: item for item in session.exec(
|
|||
|
|
select(KnowledgeFile).where(KnowledgeFile.kb_id == kb.id)
|
|||
|
|
).all()}
|
|||
|
|
for path in files:
|
|||
|
|
digest = dedup.file_hash(path)
|
|||
|
|
record = known.get(path.name)
|
|||
|
|
if record is None:
|
|||
|
|
record = KnowledgeFile(kb_id=kb.id, file_name=path.name,
|
|||
|
|
content_hash=digest, revision=1)
|
|||
|
|
elif not record.content_hash:
|
|||
|
|
# 旧数据首次补指纹时不把未知历史误报为内容变更。
|
|||
|
|
record.content_hash = digest
|
|||
|
|
record.revision = max(1, int(record.revision or 1))
|
|||
|
|
elif record.content_hash != digest:
|
|||
|
|
record.content_hash = digest
|
|||
|
|
record.revision = max(1, int(record.revision or 1)) + 1
|
|||
|
|
session.add(record)
|
|||
|
|
try:
|
|||
|
|
chunk_count = store.count(kb.id)
|
|||
|
|
except Exception:
|
|||
|
|
chunk_count = 0
|
|||
|
|
registered_count = len(known) or len(session.exec(
|
|||
|
|
select(KnowledgeFile).where(KnowledgeFile.kb_id == kb.id)
|
|||
|
|
).all())
|
|||
|
|
kb.file_count = len(files)
|
|||
|
|
kb.chunk_count = chunk_count
|
|||
|
|
if backend.is_disabled():
|
|||
|
|
kb.index_status = "disabled" if files else "empty"
|
|||
|
|
elif files and chunk_count == 0:
|
|||
|
|
kb.index_status = "needs_reindex"
|
|||
|
|
elif not files and registered_count:
|
|||
|
|
kb.index_status = "error"
|
|||
|
|
elif not files:
|
|||
|
|
kb.index_status = "empty"
|
|||
|
|
else:
|
|||
|
|
kb.index_status = "ready"
|
|||
|
|
session.add(kb)
|
|||
|
|
repaired.append({"kb_id": kb.id, "files": len(files), "chunks": chunk_count,
|
|||
|
|
"needs_reindex": bool(files and chunk_count == 0)})
|
|||
|
|
session.commit()
|
|||
|
|
return {"kbs": repaired}
|
|||
|
|
|
|||
|
|
def category_path(session: Session, category_id: int | None) -> str:
|
|||
|
|
"""把分类节点解析成“父级 / 子级”路径，写入每个向量的元数据。"""
|
|||
|
|
if not category_id:
|
|||
|
|
return "未分类"
|
|||
|
|
names = []
|
|||
|
|
seen = set()
|
|||
|
|
current = session.get(KnowledgeCategory, category_id)
|
|||
|
|
while current and current.id not in seen:
|
|||
|
|
seen.add(current.id)
|
|||
|
|
names.append(current.name)
|
|||
|
|
current = session.get(KnowledgeCategory, current.parent_id) if current.parent_id else None
|
|||
|
|
return " / ".join(reversed(names)) or "未分类"
|
|||
|
|
|
|||
|
|
def _parse_source(path: Path, kb: KnowledgeBase, progress_cb=None) -> tuple[list[dict], int]:
|
|||
|
|
if path.suffix.lower() in MEDIA_SUFFIXES:
|
|||
|
|
from . import media as media_mod
|
|||
|
|
if progress_cb:
|
|||
|
|
progress_cb(f"转写 {path.name}（本地 whisper，长音频需几分钟）", -1)
|
|||
|
|
return media_mod.transcribe_media(path), 2200
|
|||
|
|
configured = int(getattr(kb, "chunk_target", 0) or 0)
|
|||
|
|
target = configured or (2200 if path.suffix.lower() == ".srt" else chunk.TARGET)
|
|||
|
|
return parse.parse_file(path), target
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _indexed_sources(session: Session, kb: KnowledgeBase, source: Path, cache) -> list[dict]:
|
|||
|
|
"""只将原文件和索引版本一致、且实际有块的资料作为跳过上传的依据。"""
|
|||
|
|
counts = store.count_by_file(kb.id)
|
|||
|
|
if not counts:
|
|||
|
|
return []
|
|||
|
|
first_meta, legacy_texts, versions = {}, {}, {}
|
|||
|
|
for item in store.iter_chunks(kb.id):
|
|||
|
|
meta = item.get("meta") or {}
|
|||
|
|
name = meta.get("file")
|
|||
|
|
if not name:
|
|||
|
|
continue
|
|||
|
|
first_meta.setdefault(name, meta)
|
|||
|
|
versions.setdefault(name, set()).add(str(meta.get("file_hash") or ""))
|
|||
|
|
if not meta.get("file_hash") or not meta.get("file_chunk_count"):
|
|||
|
|
legacy_texts.setdefault(name, []).append(
|
|||
|
|
(int(meta.get("seq", len(legacy_texts.get(name, [])))),
|
|||
|
|
(dedup.normalized_text(meta.get("heading", "")), dedup.normalized_text(item.get("text", "")))))
|
|||
|
|
candidates = []
|
|||
|
|
for record in session.exec(select(KnowledgeFile).where(KnowledgeFile.kb_id == kb.id)).all():
|
|||
|
|
path = source / record.file_name
|
|||
|
|
if not counts.get(record.file_name) or not path.is_file():
|
|||
|
|
continue
|
|||
|
|
raw_hash = dedup.file_hash(path)
|
|||
|
|
meta = first_meta.get(record.file_name, {})
|
|||
|
|
if meta.get("file_hash") and versions.get(record.file_name) != {raw_hash}:
|
|||
|
|
continue
|
|||
|
|
if meta.get("file_chunk_count") and int(meta["file_chunk_count"]) != counts.get(record.file_name):
|
|||
|
|
continue
|
|||
|
|
body_hash = cache.get(raw_hash)
|
|||
|
|
if not meta.get("file_hash") or not meta.get("file_chunk_count"):
|
|||
|
|
# 没有版本指纹的旧索引必须与当前原文逐块核对，不能只看文件表。
|
|||
|
|
if path.suffix.lower() in MEDIA_SUFFIXES:
|
|||
|
|
continue
|
|||
|
|
units, target = _parse_source(path, kb)
|
|||
|
|
raw_chunks = chunk.chunk_texts(units, target=target)
|
|||
|
|
unique, _ = dedup.unique_chunks(raw_chunks)
|
|||
|
|
def canonical(items):
|
|||
|
|
return [(dedup.normalized_text((item.get("meta") or {}).get("heading", "")),
|
|||
|
|
dedup.normalized_text(item["text"])) for item in items]
|
|||
|
|
expected = canonical(raw_chunks)
|
|||
|
|
indexed = [text for _, text in sorted(legacy_texts.get(record.file_name, []))]
|
|||
|
|
if expected != indexed and canonical(unique) != indexed:
|
|||
|
|
continue
|
|||
|
|
body_hash = dedup.document_hash(units)
|
|||
|
|
cache.set(raw_hash, body_hash)
|
|||
|
|
candidates.append({"file_name": record.file_name, "raw_hash": raw_hash,
|
|||
|
|
"body_hash": body_hash, "path": path})
|
|||
|
|
return sorted(candidates, key=lambda item: item["file_name"])
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _duplicate_of(candidates: list[dict], raw_hash: str, body_hash: str = "") -> str:
|
|||
|
|
for candidate in candidates:
|
|||
|
|
if candidate["raw_hash"] == raw_hash or (body_hash and candidate["body_hash"] == body_hash):
|
|||
|
|
return candidate["file_name"]
|
|||
|
|
return ""
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _fill_body_hashes(candidates: list[dict], kb: KnowledgeBase, cache):
|
|||
|
|
for candidate in candidates:
|
|||
|
|
if candidate["body_hash"]:
|
|||
|
|
continue
|
|||
|
|
# 相同音视频先按文件字节去重；不同媒体不为去重而重复转写。
|
|||
|
|
if candidate["path"].suffix.lower() in MEDIA_SUFFIXES:
|
|||
|
|
continue
|
|||
|
|
units, _ = _parse_source(candidate["path"], kb)
|
|||
|
|
candidate["body_hash"] = dedup.document_hash(units)
|
|||
|
|
cache.set(candidate["raw_hash"], candidate["body_hash"])
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _store_usage(session, kb, path, file_hash, revision, chunks, operation):
|
|||
|
|
record_usage_event(
|
|||
|
|
session, idempotency_key=f"storage:kb:{kb.id}:{file_hash}",
|
|||
|
|
user_id=kb.owner_user_id, organization_id=kb.organization_id,
|
|||
|
|
metadata={"kb_id": kb.id, "file_name": path.name, "kind": "knowledge_file"},
|
|||
|
|
service_type="storage", unit="byte", quantity=path.stat().st_size, status="measured")
|
|||
|
|
record_usage_event(
|
|||
|
|
session, idempotency_key=f"embedding:index:{kb.id}:{file_hash}:{revision}:{operation}",
|
|||
|
|
user_id=kb.owner_user_id, organization_id=kb.organization_id,
|
|||
|
|
model=getattr(embed, "MODEL_NAME", "local-embedding"), service_type="embedding",
|
|||
|
|
unit="token", quantity=sum(estimate_tokens(item["text"]) for item in chunks),
|
|||
|
|
status="estimated", metadata={"kb_id": kb.id, "file_name": path.name,
|
|||
|
|
"chunk_count": len(chunks), "operation": operation})
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _import_kb_files_locked(session: Session, kb_id: int, file_paths: list[str], progress_cb=None,
|
|||
|
|
*, operation="import", candidates=None) -> dict:
|
|||
|
|
kb = session.get(KnowledgeBase, kb_id)
|
|||
|
|
if not kb:
|
|||
|
|
raise ValueError("知识库不存在")
|
|||
|
|
source = _source_dir(kb)
|
|||
|
|
source.mkdir(parents=True, exist_ok=True)
|
|||
|
|
if backend.is_disabled():
|
|||
|
|
# Lightweight mode: preserve source files and registry rows, but never
|
|||
|
|
# download/load embedding models or build a local vector index.
|
|||
|
|
imported = []
|
|||
|
|
for fp in file_paths:
|
|||
|
|
incoming = Path(fp)
|
|||
|
|
if not incoming.is_file() or incoming.suffix.lower() not in SUPPORTED_SUFFIXES | MEDIA_SUFFIXES:
|
|||
|
|
continue
|
|||
|
|
dest = source / incoming.name
|
|||
|
|
if incoming.resolve() != dest.resolve():
|
|||
|
|
shutil.copy2(incoming, dest)
|
|||
|
|
digest = dedup.file_hash(dest)
|
|||
|
|
record = session.exec(select(KnowledgeFile).where(
|
|||
|
|
KnowledgeFile.kb_id == kb_id, KnowledgeFile.file_name == dest.name)).first()
|
|||
|
|
if record is None:
|
|||
|
|
record = KnowledgeFile(kb_id=kb_id, file_name=dest.name)
|
|||
|
|
record.content_hash = digest
|
|||
|
|
record.revision = max(1, int(record.revision or 1))
|
|||
|
|
session.add(record)
|
|||
|
|
imported.append(dest.name)
|
|||
|
|
kb.file_count = len(_disk_files(source))
|
|||
|
|
kb.source_dir = str(source)
|
|||
|
|
kb.chunk_count = 0
|
|||
|
|
kb.index_status = "disabled"
|
|||
|
|
session.add(kb)
|
|||
|
|
session.commit()
|
|||
|
|
return {"files": kb.file_count, "chunks": 0, "imported_files": len(imported),
|
|||
|
|
"imported_file_names": imported, "skipped_duplicates": 0,
|
|||
|
|
"duplicate_files": [], "removed_duplicate_chunks": 0,
|
|||
|
|
"vector_disabled": True}
|
|||
|
|
cache = dedup.FingerprintCache(kb_id)
|
|||
|
|
if candidates is None:
|
|||
|
|
candidates = _indexed_sources(session, kb, source, cache)
|
|||
|
|
previous_status = kb.index_status
|
|||
|
|
kb.source_dir = str(source)
|
|||
|
|
kb.index_status = "indexing"
|
|||
|
|
session.add(kb)
|
|||
|
|
session.commit()
|
|||
|
|
result = {"files": 0, "chunks": 0, "imported_files": 0,
|
|||
|
|
"imported_file_names": [], "skipped_duplicates": 0, "duplicate_files": [],
|
|||
|
|
"removed_duplicate_chunks": 0}
|
|||
|
|
try:
|
|||
|
|
for fp in file_paths:
|
|||
|
|
incoming = Path(fp)
|
|||
|
|
if not incoming.is_file() or incoming.suffix.lower() not in SUPPORTED_SUFFIXES | MEDIA_SUFFIXES:
|
|||
|
|
continue
|
|||
|
|
file_name, dest = incoming.name, source / incoming.name
|
|||
|
|
if progress_cb:
|
|||
|
|
progress_cb(f"检查 {file_name} 是否重复", -1)
|
|||
|
|
raw_hash = dedup.file_hash(incoming)
|
|||
|
|
duplicate = _duplicate_of(candidates, raw_hash)
|
|||
|
|
units, target, body_hash = None, 0, ""
|
|||
|
|
if not duplicate:
|
|||
|
|
units, target = _parse_source(incoming, kb, progress_cb)
|
|||
|
|
body_hash = dedup.document_hash(units)
|
|||
|
|
_fill_body_hashes(candidates, kb, cache)
|
|||
|
|
duplicate = _duplicate_of(candidates, raw_hash, body_hash)
|
|||
|
|
if duplicate:
|
|||
|
|
result["skipped_duplicates"] += 1
|
|||
|
|
result["duplicate_files"].append({"file_name": file_name, "duplicate_of": duplicate})
|
|||
|
|
if progress_cb:
|
|||
|
|
progress_cb(f"{file_name}：与 {duplicate} 内容重复，已跳过", 0)
|
|||
|
|
continue
|
|||
|
|
raw_chunks = chunk.chunk_texts(units, target=target)
|
|||
|
|
chunks, removed = dedup.unique_chunks(raw_chunks)
|
|||
|
|
if not chunks:
|
|||
|
|
raise ValueError(f"{file_name} 未解析出可索引正文，已保留原文件与原索引；扫描件请先 OCR，旧版 doc 请转换为 docx")
|
|||
|
|
record = session.exec(select(KnowledgeFile).where(
|
|||
|
|
KnowledgeFile.kb_id == kb_id, KnowledgeFile.file_name == file_name)).first()
|
|||
|
|
revision = max(1, int(record.revision or 1)) if record else 1
|
|||
|
|
if record and record.content_hash != raw_hash:
|
|||
|
|
revision += 1
|
|||
|
|
for item in chunks:
|
|||
|
|
item["meta"].update(file_hash=raw_hash, file_revision=str(revision),
|
|||
|
|
document_hash=body_hash, dedup_version=1, file_chunk_count=len(chunks))
|
|||
|
|
if progress_cb:
|
|||
|
|
progress_cb(f"{file_name}：去重后 {len(chunks)} 个片段，正在建立索引" +
|
|||
|
|
(f"（跳过 {removed} 个重复片段）" if removed else ""), -1)
|
|||
|
|
# 完成解析及嵌入后才替换索引/源文件，失败不会先破坏用户旧文件。
|
|||
|
|
vectors = embed.embed_texts([item["text"] for item in chunks])
|
|||
|
|
stage = source / f".{uuid.uuid4().hex}.upload"
|
|||
|
|
backup = source / f".{uuid.uuid4().hex}.previous"
|
|||
|
|
copying = incoming.resolve() != dest.resolve()
|
|||
|
|
snapshot = None
|
|||
|
|
index_written = source_replaced = source_backed_up = False
|
|||
|
|
try:
|
|||
|
|
if copying:
|
|||
|
|
shutil.copy2(incoming, stage)
|
|||
|
|
if dedup.file_hash(stage) != raw_hash:
|
|||
|
|
raise RuntimeError("上传文件在处理过程中发生变化，请重试")
|
|||
|
|
snapshot = store.snapshot_file(kb_id, file_name)
|
|||
|
|
store.replace_file_chunks(kb_id, chunks, vectors, file_name,
|
|||
|
|
category_path(session, (record.category_id or record.ai_category_id) if record else None))
|
|||
|
|
index_written = True
|
|||
|
|
if copying:
|
|||
|
|
if dest.exists():
|
|||
|
|
dest.replace(backup)
|
|||
|
|
source_backed_up = True
|
|||
|
|
stage.replace(dest)
|
|||
|
|
source_replaced = True
|
|||
|
|
if record is None:
|
|||
|
|
record = KnowledgeFile(kb_id=kb_id, file_name=file_name)
|
|||
|
|
record.content_hash, record.revision = raw_hash, revision
|
|||
|
|
session.add(record)
|
|||
|
|
session.commit()
|
|||
|
|
except Exception as exc:
|
|||
|
|
session.rollback()
|
|||
|
|
rollback_errors = []
|
|||
|
|
if index_written and snapshot is not None:
|
|||
|
|
try:
|
|||
|
|
store.restore_file(kb_id, file_name, snapshot)
|
|||
|
|
except Exception as restore_exc:
|
|||
|
|
rollback_errors.append(f"索引恢复失败：{restore_exc}")
|
|||
|
|
if copying:
|
|||
|
|
try:
|
|||
|
|
if source_replaced:
|
|||
|
|
dest.unlink(missing_ok=True)
|
|||
|
|
if source_backed_up:
|
|||
|
|
backup.replace(dest)
|
|||
|
|
except OSError as restore_exc:
|
|||
|
|
rollback_errors.append(f"源文件恢复失败（备份 {backup.name}）：{restore_exc}")
|
|||
|
|
if rollback_errors:
|
|||
|
|
raise RuntimeError(f"导入失败：{exc}；{'；'.join(rollback_errors)}") from exc
|
|||
|
|
raise
|
|||
|
|
finally:
|
|||
|
|
stage.unlink(missing_ok=True)
|
|||
|
|
backup.unlink(missing_ok=True)
|
|||
|
|
_store_usage(session, kb, dest, raw_hash, revision, chunks, operation)
|
|||
|
|
result["imported_files"] += 1
|
|||
|
|
result["imported_file_names"].append(file_name)
|
|||
|
|
result["chunks"] += len(chunks)
|
|||
|
|
result["removed_duplicate_chunks"] += removed
|
|||
|
|
candidates[:] = [candidate for candidate in candidates if candidate["file_name"] != file_name]
|
|||
|
|
candidates.append({"file_name": file_name, "raw_hash": raw_hash,
|
|||
|
|
"body_hash": body_hash, "path": dest})
|
|||
|
|
cache.set(raw_hash, body_hash)
|
|||
|
|
if progress_cb:
|
|||
|
|
progress_cb(f"{file_name}：已导入 {len(chunks)} 个片段" +
|
|||
|
|
(f"，去重 {removed} 个片段" if removed else ""), len(chunks))
|
|||
|
|
kb.file_count = len(_disk_files(source))
|
|||
|
|
kb.chunk_count = store.count(kb_id)
|
|||
|
|
kb.index_status = "ready" if kb.chunk_count else "empty"
|
|||
|
|
session.add(kb)
|
|||
|
|
session.commit()
|
|||
|
|
cache.save()
|
|||
|
|
result["files"] = kb.file_count
|
|||
|
|
_schedule_lexical_prewarm(kb.id, int(getattr(kb, "organization_id", 1) or 1))
|
|||
|
|
return result
|
|||
|
|
except Exception:
|
|||
|
|
session.rollback()
|
|||
|
|
kb = session.get(KnowledgeBase, kb_id)
|
|||
|
|
kb.file_count = len(_disk_files(source))
|
|||
|
|
kb.chunk_count = store.count(kb_id)
|
|||
|
|
kb.index_status = previous_status if kb.chunk_count else "error"
|
|||
|
|
session.add(kb)
|
|||
|
|
session.commit()
|
|||
|
|
cache.save()
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
|
|||
|
|
def import_kb_files(session: Session, kb_id: int, file_paths: list[str], progress_cb=None) -> dict:
|
|||
|
|
"""同库文件/正文精确去重，再安全增量导入；改名副本无需再次嵌入。"""
|
|||
|
|
with dedup.library_lock(kb_id):
|
|||
|
|
result = _import_kb_files_locked(session, kb_id, file_paths, progress_cb)
|
|||
|
|
# The upload transaction is committed before this call. Build the full
|
|||
|
|
# lexical index asynchronously so the first user search is not the one
|
|||
|
|
# paying the index-construction cost.
|
|||
|
|
kb = session.get(KnowledgeBase, int(kb_id))
|
|||
|
|
if kb and kb.index_status == "ready" and kb.chunk_count:
|
|||
|
|
hybrid.schedule_prewarm(kb.organization_id, kb.id)
|
|||
|
|
# 新资料入库后后台自动做一次主题无关的全文阅读，写入可复用缓存。
|
|||
|
|
# 任务独立于上传请求，模型不可用或阅读失败都不会影响索引结果。
|
|||
|
|
imported = result.get("imported_file_names") or []
|
|||
|
|
if imported:
|
|||
|
|
try:
|
|||
|
|
from ..training_reading import schedule_auto_read
|
|||
|
|
schedule_auto_read(session, kb.id, imported)
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
|
|||
|
|
def rebuild_kb(session: Session, kb_id: int, progress_cb=None) -> dict:
|
|||
|
|
with dedup.library_lock(kb_id):
|
|||
|
|
result = _rebuild_kb_locked(session, kb_id, progress_cb)
|
|||
|
|
kb = session.get(KnowledgeBase, int(kb_id))
|
|||
|
|
if kb and kb.index_status == "ready" and kb.chunk_count:
|
|||
|
|
hybrid.schedule_prewarm(kb.organization_id, kb.id)
|
|||
|
|
# A rebuild may discover new or changed files outside the regular upload
|
|||
|
|
# route. Queue the same topic independent reading pass so later training
|
|||
|
|
# map assessments can reuse the saved summaries immediately.
|
|||
|
|
imported = result.get("imported_file_names") or []
|
|||
|
|
if imported:
|
|||
|
|
try:
|
|||
|
|
from ..training_reading import schedule_auto_read
|
|||
|
|
schedule_auto_read(session, kb.id, imported)
|
|||
|
|
except Exception:
|
|||
|
|
# Index rebuild succeeded; an unavailable model must not turn it
|
|||
|
|
# into a failed request. The next upload/rebuild can retry reading.
|
|||
|
|
pass
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
|
|||
|
|
def prewarm_ready_indexes(session: Session) -> int:
|
|||
|
|
"""Queue lexical indexes and (when enabled) the initial generic reading pass.
|
|||
|
|
|
|||
|
|
The initial AI pass is opt-in at process level because an existing large
|
|||
|
|
installation can contain many hours of material. Uploads and rebuilds
|
|||
|
|
always schedule their changed files independently below. Failed or
|
|||
|
|
interrupted passes are recovered on startup regardless of that opt-in.
|
|||
|
|
"""
|
|||
|
|
queued = 0
|
|||
|
|
bootstrap = os.environ.get("YAWAN_AUTO_READ_EXISTING", "").strip().lower() in {"1", "true", "yes", "on"}
|
|||
|
|
from ..training_reading import read_status, schedule_auto_read, auto_read_paused, _write_auto_status
|
|||
|
|
now_epoch = time.time()
|
|||
|
|
for kb in session.exec(select(KnowledgeBase).where(
|
|||
|
|
KnowledgeBase.index_status == "ready", KnowledgeBase.chunk_count > 0)).all():
|
|||
|
|
if hybrid.schedule_prewarm(kb.organization_id, kb.id):
|
|||
|
|
queued += 1
|
|||
|
|
# Existing installations opt in to the initial full pass. A failed
|
|||
|
|
# pass or a worker interrupted by process restart is always recovered,
|
|||
|
|
# even when the opt-in is absent, so a stale status cannot strand a
|
|||
|
|
# library forever. ``updated_epoch`` is written atomically by the
|
|||
|
|
# reader and a generous timeout avoids duplicating a legitimately long
|
|||
|
|
# batch still running in another process.
|
|||
|
|
status = read_status(kb)
|
|||
|
|
if auto_read_paused() and status.get("status") in {"reading", "queued"}:
|
|||
|
|
_write_auto_status(kb.organization_id or 1, kb.id,
|
|||
|
|
{**status, "status": "paused", "error": "管理员已暂停自动预读"})
|
|||
|
|
continue
|
|||
|
|
try:
|
|||
|
|
updated = float(status.get("updated_epoch") or 0)
|
|||
|
|
except (TypeError, ValueError):
|
|||
|
|
updated = 0
|
|||
|
|
stale = (status.get("status") == "reading" and
|
|||
|
|
now_epoch - updated > 900)
|
|||
|
|
retry = status.get("status") in {"error", "queued"} or stale
|
|||
|
|
if bootstrap or retry:
|
|||
|
|
try:
|
|||
|
|
names = [item.file_name for item in session.exec(
|
|||
|
|
select(KnowledgeFile).where(KnowledgeFile.kb_id == kb.id)
|
|||
|
|
).all() if item.file_name]
|
|||
|
|
if names:
|
|||
|
|
schedule_auto_read(session, kb.id, names)
|
|||
|
|
except Exception:
|
|||
|
|
# A model or cache problem must never prevent the API from
|
|||
|
|
# starting; the next rebuild/upload can retry the reading pass.
|
|||
|
|
pass
|
|||
|
|
return queued
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _rebuild_kb_locked(session: Session, kb_id: int, progress_cb=None) -> dict:
|
|||
|
|
kb = session.get(KnowledgeBase, kb_id)
|
|||
|
|
if not kb:
|
|||
|
|
raise ValueError("知识库不存在")
|
|||
|
|
source = _source_dir(kb)
|
|||
|
|
previous = {"file_count": kb.file_count, "chunk_count": kb.chunk_count, "index_status": kb.index_status}
|
|||
|
|
previous_records = {record.file_name: (record.content_hash, record.revision)
|
|||
|
|
for record in session.exec(select(KnowledgeFile).where(KnowledgeFile.kb_id == kb_id)).all()}
|
|||
|
|
try:
|
|||
|
|
snapshot = store.snapshot_kb(kb_id)
|
|||
|
|
except Exception as exc:
|
|||
|
|
if int(kb.chunk_count or 0) > 0:
|
|||
|
|
raise RuntimeError(f"无法创建索引快照，已取消重建以保护原数据：{str(exc)[:300]}") from exc
|
|||
|
|
snapshot = {"ids": [], "documents": [], "embeddings": [], "metadatas": []}
|
|||
|
|
# 源目录丢失属于异常，不能把它解释成用户要删除全部索引。
|
|||
|
|
if not source.exists() and int(kb.chunk_count or 0) > 0:
|
|||
|
|
raise ValueError("知识库源目录不存在，已取消重建并保留原索引")
|
|||
|
|
files = [str(path) for path in _disk_files(source)]
|
|||
|
|
try:
|
|||
|
|
# 保留旧索引直到对应文件解析/嵌入完成；重建也应用同库文件去重。
|
|||
|
|
result = _import_kb_files_locked(session, kb_id, files, progress_cb,
|
|||
|
|
operation="rebuild", candidates=[])
|
|||
|
|
canonical = set(result["imported_file_names"])
|
|||
|
|
previous_names = {meta.get("file") for meta in snapshot.get("metadatas", []) if meta}
|
|||
|
|
for file_name in previous_names - canonical:
|
|||
|
|
if file_name:
|
|||
|
|
store.delete_file(kb_id, file_name)
|
|||
|
|
kb.file_count = len(_disk_files(source))
|
|||
|
|
kb.chunk_count = store.count(kb_id)
|
|||
|
|
kb.index_status = "ready" if kb.chunk_count else "empty"
|
|||
|
|
session.add(kb)
|
|||
|
|
session.commit()
|
|||
|
|
return {**result, "files": kb.file_count, "rolled_back": False}
|
|||
|
|
except Exception as exc:
|
|||
|
|
session.rollback()
|
|||
|
|
restore_error = ""
|
|||
|
|
try:
|
|||
|
|
store.restore_kb(kb_id, snapshot)
|
|||
|
|
for record in session.exec(select(KnowledgeFile).where(KnowledgeFile.kb_id == kb_id)).all():
|
|||
|
|
if record.file_name in previous_records:
|
|||
|
|
record.content_hash, record.revision = previous_records[record.file_name]
|
|||
|
|
session.add(record)
|
|||
|
|
else:
|
|||
|
|
session.delete(record)
|
|||
|
|
kb.file_count = previous["file_count"]
|
|||
|
|
kb.chunk_count = previous["chunk_count"]
|
|||
|
|
kb.index_status = previous["index_status"]
|
|||
|
|
except Exception as restore_exc:
|
|||
|
|
kb.index_status = "error"
|
|||
|
|
restore_error = f"；回滚失败：{str(restore_exc)[:180]}"
|
|||
|
|
session.add(kb)
|
|||
|
|
session.commit()
|
|||
|
|
raise RuntimeError(f"知识库重建失败，已尝试恢复原索引：{str(exc)[:300]}{restore_error}") from exc
|
|||
|
|
|
|||
|
|
def consumed_chunk_ids(session: Session, agent_id: int, kb_ids: list[int]) -> dict[int, set[str]]:
|
|||
|
|
"""返回仅对当前 Agent 不可再次使用的知识片段。"""
|
|||
|
|
if not agent_id or not kb_ids:
|
|||
|
|
return {}
|
|||
|
|
usages = session.exec(select(KnowledgeUsage).where(
|
|||
|
|
KnowledgeUsage.agent_id == agent_id,
|
|||
|
|
or_(KnowledgeUsage.status.in_(["reserved", "consumed"]), KnowledgeUsage.banned == True), # noqa: E712
|
|||
|
|
)).all()
|
|||
|
|
result: dict[int, set[str]] = {kid: set() for kid in kb_ids}
|
|||
|
|
for usage in usages:
|
|||
|
|
if usage.kb_id in result:
|
|||
|
|
result[usage.kb_id].add(usage.chunk_id)
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
|
|||
|
|
def reserve_chunks(session: Session, agent_id: int, campaign_id: int,
|
|||
|
|
segment_id: int, hits: list[dict]) -> list[dict]:
|
|||
|
|
"""模型调用前原子预留；正文落盘成功后才转为正式 consumed。"""
|
|||
|
|
# 预取候选状态并使用 executemany，避免每个 chunk 两次往返 SQLite。
|
|||
|
|
# 这条路径通常一次处理几十个片段，批量写入可显著减少写锁持有时间。
|
|||
|
|
unique_hits, seen = [], set()
|
|||
|
|
for hit in hits:
|
|||
|
|
kb_id = int(hit.get("kb_id") or 0)
|
|||
|
|
chunk_id = str(hit.get("id") or "")
|
|||
|
|
if not kb_id or not chunk_id:
|
|||
|
|
continue
|
|||
|
|
key = (kb_id, chunk_id)
|
|||
|
|
if key in seen:
|
|||
|
|
continue
|
|||
|
|
seen.add(key)
|
|||
|
|
unique_hits.append(hit)
|
|||
|
|
if not unique_hits or not agent_id:
|
|||
|
|
return []
|
|||
|
|
kb_ids = sorted({int(hit.get("kb_id") or 0) for hit in unique_hits})
|
|||
|
|
existing = session.exec(select(KnowledgeUsage).where(
|
|||
|
|
KnowledgeUsage.agent_id == int(agent_id),
|
|||
|
|
KnowledgeUsage.kb_id.in_(kb_ids),
|
|||
|
|
KnowledgeUsage.chunk_id.in_([str(hit.get("id") or "") for hit in unique_hits])
|
|||
|
|
)).all()
|
|||
|
|
state = {(int(row.kb_id), str(row.chunk_id)): row for row in existing}
|
|||
|
|
inserts, reopens, candidates = [], [], []
|
|||
|
|
for hit in unique_hits:
|
|||
|
|
kb_id = int(hit.get("kb_id") or 0)
|
|||
|
|
chunk_id = str(hit.get("id") or "")
|
|||
|
|
values = {
|
|||
|
|
"agent_id": agent_id, "kb_id": kb_id, "chunk_id": chunk_id,
|
|||
|
|
"file_name": str(hit.get("meta", {}).get("file") or ""),
|
|||
|
|
"heading": str(hit.get("meta", {}).get("heading") or ""),
|
|||
|
|
"excerpt": str(hit.get("text") or "")[:1000],
|
|||
|
|
"campaign_id": campaign_id, "segment_id": segment_id,
|
|||
|
|
"used_at": now(),
|
|||
|
|
}
|
|||
|
|
row = state.get((kb_id, chunk_id))
|
|||
|
|
if row is None:
|
|||
|
|
inserts.append(values); candidates.append(hit)
|
|||
|
|
elif str(row.status) == "available" and not bool(row.banned):
|
|||
|
|
reopens.append(values); candidates.append(hit)
|
|||
|
|
if inserts:
|
|||
|
|
session.execute(text("""
|
|||
|
|
INSERT OR IGNORE INTO knowledgeusage
|
|||
|
|
(agent_id, kb_id, chunk_id, file_name, heading, excerpt, campaign_id, segment_id, status, used_at, released_at, banned)
|
|||
|
|
VALUES (:agent_id, :kb_id, :chunk_id, :file_name, :heading, :excerpt, :campaign_id, :segment_id, 'reserved', :used_at, '', 0)
|
|||
|
|
"""), inserts)
|
|||
|
|
if reopens:
|
|||
|
|
session.execute(text("""
|
|||
|
|
UPDATE knowledgeusage
|
|||
|
|
SET file_name = :file_name, heading = :heading, excerpt = :excerpt,
|
|||
|
|
campaign_id = :campaign_id, segment_id = :segment_id,
|
|||
|
|
status = 'reserved', used_at = :used_at, released_at = ''
|
|||
|
|
WHERE agent_id = :agent_id AND kb_id = :kb_id AND chunk_id = :chunk_id
|
|||
|
|
AND status = 'available' AND banned = 0
|
|||
|
|
"""), reopens)
|
|||
|
|
# 并发请求可能在预取后先认领了同一片段；以最终状态确认本事务实际
|
|||
|
|
# 拿到的键，避免把 INSERT OR IGNORE 的冲突行误报为已预留。
|
|||
|
|
claimed_keys = set()
|
|||
|
|
if candidates:
|
|||
|
|
rows = session.exec(select(KnowledgeUsage).where(
|
|||
|
|
KnowledgeUsage.agent_id == int(agent_id),
|
|||
|
|
KnowledgeUsage.kb_id.in_(kb_ids),
|
|||
|
|
KnowledgeUsage.chunk_id.in_([str(hit.get("id") or "") for hit in candidates]),
|
|||
|
|
KnowledgeUsage.status == "reserved",
|
|||
|
|
KnowledgeUsage.campaign_id == int(campaign_id),
|
|||
|
|
KnowledgeUsage.segment_id == int(segment_id),
|
|||
|
|
)).all()
|
|||
|
|
claimed_keys = {(int(row.kb_id), str(row.chunk_id)) for row in rows}
|
|||
|
|
session.commit()
|
|||
|
|
return [hit for hit in candidates if (int(hit.get("kb_id") or 0), str(hit.get("id") or "")) in claimed_keys]
|
|||
|
|
|
|||
|
|
def consume_reserved_chunks(session: Session, agent_id: int, campaign_id: int,
|
|||
|
|
segment_id: int, hits: list[dict]) -> None:
|
|||
|
|
"""只有正文和锚点成功写入后才视为当前 Agent 已消费。"""
|
|||
|
|
params = [{"agent_id": agent_id, "kb_id": int(hit.get("kb_id") or 0),
|
|||
|
|
"chunk_id": str(hit.get("id") or ""), "campaign_id": campaign_id,
|
|||
|
|
"segment_id": segment_id} for hit in hits if hit.get("kb_id") and hit.get("id")]
|
|||
|
|
if params:
|
|||
|
|
session.execute(text("""
|
|||
|
|
UPDATE knowledgeusage SET status = 'consumed'
|
|||
|
|
WHERE agent_id = :agent_id AND kb_id = :kb_id AND chunk_id = :chunk_id
|
|||
|
|
AND campaign_id = :campaign_id AND segment_id = :segment_id AND status = 'reserved'
|
|||
|
|
"""), params)
|
|||
|
|
session.commit()
|
|||
|
|
|
|||
|
|
def release_reserved_chunks(session: Session, agent_id: int, campaign_id: int,
|
|||
|
|
segment_id: int, hits: list[dict]) -> None:
|
|||
|
|
"""失败、重试或锚点不合格时归还未落盘素材。"""
|
|||
|
|
params = [{"agent_id": agent_id, "kb_id": int(hit.get("kb_id") or 0),
|
|||
|
|
"chunk_id": str(hit.get("id") or ""), "campaign_id": campaign_id,
|
|||
|
|
"segment_id": segment_id, "released_at": now()} for hit in hits if hit.get("kb_id") and hit.get("id")]
|
|||
|
|
if params:
|
|||
|
|
session.execute(text("""
|
|||
|
|
UPDATE knowledgeusage SET status = 'available', released_at = :released_at
|
|||
|
|
WHERE agent_id = :agent_id AND kb_id = :kb_id AND chunk_id = :chunk_id
|
|||
|
|
AND campaign_id = :campaign_id AND segment_id = :segment_id AND status = 'reserved'
|
|||
|
|
"""), params)
|
|||
|
|
session.commit()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def search_kbs_remote(session: Session, kb_ids: list[int], query: str, k: int = 6,
|
|||
|
|
excluded_chunk_ids: dict[int, set[str]] | None = None, *,
|
|||
|
|
campaign_id: int = 0, trace_id: str = "", user_id: int = 0,
|
|||
|
|
organization_id: int = 1, category_ids: list[int] | None = None,
|
|||
|
|
file_whitelists: dict[int, set[str]] | None = None) -> list[dict]:
|
|||
|
|
"""有界进程池检索；分类和调用方文件范围在召回前按知识库取交集。"""
|
|||
|
|
kids = list(dict.fromkeys(int(kid) for kid in kb_ids))
|
|||
|
|
if not kids or not str(query or '').strip() or k <= 0:
|
|||
|
|
return []
|
|||
|
|
if backend.is_disabled():
|
|||
|
|
return _lightweight_keyword_search(
|
|||
|
|
session, kids, query, k=k, excluded_chunk_ids=excluded_chunk_ids,
|
|||
|
|
campaign_id=campaign_id, trace_id=trace_id, user_id=user_id,
|
|||
|
|
organization_id=organization_id, category_ids=category_ids,
|
|||
|
|
file_whitelists=file_whitelists, validate_access=True)
|
|||
|
|
whitelists = {int(kid): set(names) for kid, names in (file_whitelists or {}).items()}
|
|||
|
|
if category_ids:
|
|||
|
|
for kid in kids:
|
|||
|
|
allowed = _category_file_whitelist(session, kid, list(category_ids))
|
|||
|
|
if allowed is not None:
|
|||
|
|
whitelists[kid] = allowed & whitelists[kid] if kid in whitelists else allowed
|
|||
|
|
# 显式空集合表示没有获选文件，绝不能当作“不限范围”。
|
|||
|
|
eligible = [kid for kid in kids if kid not in whitelists or whitelists[kid]]
|
|||
|
|
if not eligible:
|
|||
|
|
return []
|
|||
|
|
# Warm missing full-library FTS indexes from the API process. The search
|
|||
|
|
# worker remains focused on the bounded vector query and is not burdened by
|
|||
|
|
# a long background SQLite build that could be killed with the worker.
|
|||
|
|
for kid in eligible:
|
|||
|
|
try:
|
|||
|
|
rev = hybrid.revision(kid)
|
|||
|
|
if not hybrid.index_available(organization_id, kid, rev):
|
|||
|
|
hybrid.schedule_prewarm(organization_id, kid, rev)
|
|||
|
|
except Exception:
|
|||
|
|
# Prewarming is best effort; vector recall below remains available.
|
|||
|
|
pass
|
|||
|
|
hits = _search_kbs_worker(session, eligible, query, k, excluded_chunk_ids,
|
|||
|
|
campaign_id, trace_id, user_id, organization_id,
|
|||
|
|
file_whitelists=whitelists)
|
|||
|
|
# 最后仅遍历一次防御校验，不因选择多个库而重复或扩大结果。
|
|||
|
|
results, seen = [], set()
|
|||
|
|
for hit in hits:
|
|||
|
|
kid = int(hit.get('kb_id') or 0)
|
|||
|
|
if kid not in eligible:
|
|||
|
|
continue
|
|||
|
|
meta = hit.get('meta') or {}
|
|||
|
|
if kid in whitelists and meta.get('file') not in whitelists[kid]:
|
|||
|
|
continue
|
|||
|
|
chunk_id = hit.get('id')
|
|||
|
|
identity = (kid, str(chunk_id)) if chunk_id else (kid, meta.get('file'), meta.get('seq'), hit.get('text'))
|
|||
|
|
if identity not in seen:
|
|||
|
|
seen.add(identity)
|
|||
|
|
results.append(hit)
|
|||
|
|
return results
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _search_kbs_worker(session: Session, kb_ids: list[int], query: str, k: int = 6,
|
|||
|
|
excluded_chunk_ids: dict[int, set[str]] | None = None,
|
|||
|
|
campaign_id: int = 0, trace_id: str = "", user_id: int = 0,
|
|||
|
|
organization_id: int = 1, file_whitelists: dict[int, set[str]] | None = None) -> list[dict]:
|
|||
|
|
from .worker_pool import get_pool, RagWorkerError
|
|||
|
|
payload = {
|
|||
|
|
'kb_ids': [int(kid) for kid in kb_ids], 'query': query, 'k': k,
|
|||
|
|
'excluded': {str(kid): sorted(ids) for kid, ids in (excluded_chunk_ids or {}).items()},
|
|||
|
|
'campaign_id': int(campaign_id or 0), 'trace_id': trace_id or '',
|
|||
|
|
'user_id': int(user_id or 0), 'organization_id': int(organization_id or 1),
|
|||
|
|
'file_whitelists': {str(kid): sorted(names) for kid, names in (file_whitelists or {}).items()},
|
|||
|
|
}
|
|||
|
|
try:
|
|||
|
|
response = get_pool().request(payload)
|
|||
|
|
except RagWorkerError as exc:
|
|||
|
|
metric = getattr(exc, "retrieval_metric", None)
|
|||
|
|
if metric:
|
|||
|
|
perf_metrics.ingest_rag(metric)
|
|||
|
|
from .. import system_log
|
|||
|
|
import logging
|
|||
|
|
message = f'知识库检索失败 · 知识库：{kb_ids} · 原因：{system_log.redact(exc, 500)}'
|
|||
|
|
try:
|
|||
|
|
# 独立短事务记录错误，不提交调用方尚未完成的业务变更。
|
|||
|
|
with Session(session.get_bind()) as log_session:
|
|||
|
|
if log_session.get_bind().dialect.name == 'sqlite':
|
|||
|
|
log_session.connection().exec_driver_sql('PRAGMA busy_timeout=100')
|
|||
|
|
system_log.write(log_session, source='rag', message=message, trace_id=trace_id,
|
|||
|
|
status='error')
|
|||
|
|
except Exception:
|
|||
|
|
logging.getLogger('yawen.rag').error(message)
|
|||
|
|
# 不回退主进程执行：避免超时后继续无界计算并重复计量。
|
|||
|
|
raise
|
|||
|
|
timings = {**response.get('worker_timings', {}), **response.get('transport_timings', {})}
|
|||
|
|
if response.get("retrieval_metric"):
|
|||
|
|
metric = dict(response["retrieval_metric"])
|
|||
|
|
# Transport diagnostics are measured in the API process and are kept
|
|||
|
|
# alongside worker phase metrics for cold start/IPC analysis.
|
|||
|
|
# ``worker_ms`` includes the complete child search (embedding, vector
|
|||
|
|
# recall and post-processing). Using it as IPC time made the
|
|||
|
|
# diagnostics wildly overstate transport overhead. The worker pool
|
|||
|
|
# computes the residual pipe/JSON time explicitly as ``ipc_ms``;
|
|||
|
|
# retain a zero fallback for old workers that do not provide it.
|
|||
|
|
metric["embedding_ipc_ms"] = timings.get(
|
|||
|
|
"ipc_ms", metric.get("embedding_ipc_ms", 0.0)
|
|||
|
|
)
|
|||
|
|
# Keep queue wait separate from child search and IPC so a saturated
|
|||
|
|
# resident worker is distinguishable from slow disk/vector recall.
|
|||
|
|
metric["queue_wait_ms"] = timings.get("queue_ms", 0.0)
|
|||
|
|
metric["transport_total_ms"] = timings.get("total_ms", 0.0)
|
|||
|
|
metric["worker_pid"] = timings.get("worker_pid")
|
|||
|
|
metric["worker_start_ms"] = timings.get("process_start_ms", 0.0)
|
|||
|
|
metric["worker_reused"] = timings.get("process_reused", False)
|
|||
|
|
# Normalize the worker's compact embedding envelope to the telemetry
|
|||
|
|
# schema used by the API diagnostics endpoint.
|
|||
|
|
metric["embedding_model"] = metric.get("embedding_model") or metric.get("model")
|
|||
|
|
metric["embedding_version"] = metric.get("embedding_version") or metric.get("model_version")
|
|||
|
|
metric["embedding_pid"] = metric.get("embedding_pid") or metric.get("pid")
|
|||
|
|
metric["embedding_process_start_ms"] = timings.get(
|
|||
|
|
"process_start_ms", metric.get("embedding_process_start_ms", 0.0)
|
|||
|
|
)
|
|||
|
|
metric["embedding_cache_hit"] = metric.get("embedding_cache_hit", metric.get("cache_hit", False))
|
|||
|
|
metric["embedding_model_load_ms"] = metric.get("embedding_model_load_ms", metric.get("model_load_ms", 0.0))
|
|||
|
|
metric["embedding_tokenizer_ms"] = metric.get("embedding_tokenizer_ms", metric.get("tokenizer_ms", 0.0))
|
|||
|
|
metric["embedding_encode_ms"] = metric.get("embedding_encode_ms", metric.get("encode_ms", 0.0))
|
|||
|
|
metric["embedding_postprocess_ms"] = metric.get("embedding_postprocess_ms", metric.get("postprocess_ms", 0.0))
|
|||
|
|
metric["embedding_count"] = metric.get("embedding_count", 0)
|
|||
|
|
metric["embedding_total_ms"] = sum(float(metric.get(name, 0.0) or 0.0)
|
|||
|
|
for name in ("embedding_model_load_ms", "embedding_tokenizer_ms",
|
|||
|
|
"embedding_encode_ms", "embedding_postprocess_ms"))
|
|||
|
|
perf_metrics.ingest_rag(metric)
|
|||
|
|
for item in response['results']:
|
|||
|
|
item['transport_timings'] = timings
|
|||
|
|
return response['results']
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
_LIGHTWEIGHT_TEXT_SUFFIXES = SUPPORTED_SUFFIXES - {".doc"}
|
|||
|
|
_LIGHTWEIGHT_ASCII_RE = re.compile(r"[a-z0-9_][a-z0-9_.+-]*", re.IGNORECASE)
|
|||
|
|
_LIGHTWEIGHT_CJK_RE = re.compile(r"[\u4e00-\u9fff]{2,}")
|
|||
|
|
_LIGHTWEIGHT_NUMBER_RE = re.compile(r"\d+(?:[-–.]\d+)+")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _lightweight_chunk_target(path: Path, kb: KnowledgeBase | None = None) -> int:
|
|||
|
|
configured = int(getattr(kb, "chunk_target", 0) or 0)
|
|||
|
|
return configured or (2200 if path.suffix.lower() == ".srt" else chunk.TARGET)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _lightweight_chunk_id(kb_id: int, file_name: str, content_hash: str,
|
|||
|
|
heading: str, seq: int, text: str) -> str:
|
|||
|
|
return hashlib.sha256(
|
|||
|
|
f"{kb_id}\0{file_name}\0{content_hash or ''}\0{heading}\0{seq}\0{text}".encode("utf-8")
|
|||
|
|
).hexdigest()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def iter_lightweight_chunks(session: Session, kb_ids: list[int], *,
|
|||
|
|
excluded_chunk_ids: dict[int, set[str]] | None = None,
|
|||
|
|
category_ids_by_kb: dict[int, list[int]] | None = None,
|
|||
|
|
file_whitelists: dict[int, set[str]] | None = None,
|
|||
|
|
apply_quarantine: bool = True) -> list[dict]:
|
|||
|
|
"""Return all on-demand lightweight chunks with the same IDs used by search.
|
|||
|
|
|
|||
|
|
Vector-disabled mode never persists embeddings or local vectors. It parses
|
|||
|
|
registered source files at query time; planner validation uses this helper to
|
|||
|
|
inspect the exact material pool that DSH Worker will see.
|
|||
|
|
"""
|
|||
|
|
kids = list(dict.fromkeys(int(kid) for kid in kb_ids))
|
|||
|
|
exclusions = {int(kid): set(values) for kid, values in (excluded_chunk_ids or {}).items()}
|
|||
|
|
specified_ranges = {int(kid): set(names) for kid, names in (file_whitelists or {}).items()}
|
|||
|
|
out: list[dict] = []
|
|||
|
|
for kid in kids:
|
|||
|
|
kb = session.exec(select(KnowledgeBase).where(KnowledgeBase.id == kid)
|
|||
|
|
.execution_options(populate_existing=True)).first()
|
|||
|
|
if kb is None or str(kb.index_status or "") == "indexing":
|
|||
|
|
continue
|
|||
|
|
whitelist = _category_file_whitelist(
|
|||
|
|
session, kid, list((category_ids_by_kb or {}).get(kid) or []))
|
|||
|
|
if kid in specified_ranges:
|
|||
|
|
specified = specified_ranges[kid]
|
|||
|
|
whitelist = specified if whitelist is None else whitelist & specified
|
|||
|
|
quarantined = file_exclusions(session, kid) if apply_quarantine else set()
|
|||
|
|
if whitelist is not None:
|
|||
|
|
whitelist = whitelist - quarantined
|
|||
|
|
if whitelist is not None and not whitelist:
|
|||
|
|
continue
|
|||
|
|
source = _source_dir(kb)
|
|||
|
|
records = session.exec(select(KnowledgeFile).where(KnowledgeFile.kb_id == kid)).all()
|
|||
|
|
for record in sorted(records, key=lambda item: item.file_name):
|
|||
|
|
file_name = str(record.file_name or "")
|
|||
|
|
if whitelist is not None and file_name not in whitelist:
|
|||
|
|
continue
|
|||
|
|
if file_name in quarantined:
|
|||
|
|
continue
|
|||
|
|
path = source / file_name
|
|||
|
|
if not path.is_file() or path.suffix.lower() not in _LIGHTWEIGHT_TEXT_SUFFIXES:
|
|||
|
|
continue
|
|||
|
|
try:
|
|||
|
|
units = parse.parse_file(path)
|
|||
|
|
chunks = chunk.chunk_texts(units, target=_lightweight_chunk_target(path, kb))
|
|||
|
|
except Exception as exc:
|
|||
|
|
logger.warning("轻量素材库存解析失败，已跳过 %s/%s：%s", kid, file_name, exc)
|
|||
|
|
continue
|
|||
|
|
excluded = exclusions.get(kid, set())
|
|||
|
|
for seq, item in enumerate(chunks):
|
|||
|
|
body = str(item.get("text") or "").strip()
|
|||
|
|
if not body:
|
|||
|
|
continue
|
|||
|
|
meta = dict(item.get("meta") or {})
|
|||
|
|
heading = str(meta.get("heading") or "")
|
|||
|
|
chunk_id = _lightweight_chunk_id(
|
|||
|
|
kid, file_name, str(record.content_hash or ""), heading, seq, body)
|
|||
|
|
if chunk_id in excluded:
|
|||
|
|
continue
|
|||
|
|
item_meta = dict(meta)
|
|||
|
|
item_meta.update({"kb_id": kid, "file": file_name, "heading": heading,
|
|||
|
|
"seq": seq, "retrieval_mode": "keyword-light"})
|
|||
|
|
out.append({"id": chunk_id, "text": body, "meta": item_meta,
|
|||
|
|
"kb_id": kid, "distance": 0.0})
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
|
|||
|
|
def simulate_lightweight_plan(session: Session, kb_ids: list[int], queries: list[str], *,
|
|||
|
|
excluded_chunk_ids: dict[int, set[str]] | None = None,
|
|||
|
|
category_ids_by_kb: dict[int, list[int]] | None = None,
|
|||
|
|
file_whitelists: dict[int, set[str]] | None = None,
|
|||
|
|
initial_top_k: int = 6, widened_top_k: int = 24,
|
|||
|
|
char_budget: int = 12000) -> dict:
|
|||
|
|
"""Deterministically replay per-node keyword recall and chunk reservation.
|
|||
|
|
|
|||
|
|
The plan is approvable only if every planned node can reserve at least one
|
|||
|
|
source chunk after earlier nodes have consumed their selected materials. This
|
|||
|
|
catches exhausted/mismatched lightweight KBs before Worker execution.
|
|||
|
|
"""
|
|||
|
|
pool = iter_lightweight_chunks(
|
|||
|
|
session, kb_ids, excluded_chunk_ids=excluded_chunk_ids,
|
|||
|
|
category_ids_by_kb=category_ids_by_kb, file_whitelists=file_whitelists)
|
|||
|
|
available = {str(item["id"]): item for item in pool}
|
|||
|
|
node_results: list[dict] = []
|
|||
|
|
for index, query_text in enumerate(queries, start=1):
|
|||
|
|
query_text = str(query_text or "")
|
|||
|
|
terms = _lightweight_terms(query_text)
|
|||
|
|
scored: list[dict] = []
|
|||
|
|
if terms:
|
|||
|
|
for item in available.values():
|
|||
|
|
score, matched = _lightweight_score(item["text"], terms)
|
|||
|
|
if score > 0 and matched > 0:
|
|||
|
|
candidate = dict(item)
|
|||
|
|
candidate["_score"] = score
|
|||
|
|
candidate["_matched"] = matched
|
|||
|
|
scored.append(candidate)
|
|||
|
|
scored.sort(key=lambda item: (-float(item["_score"]), -int(item["_matched"]),
|
|||
|
|
int(item["kb_id"]),
|
|||
|
|
str(item["meta"].get("file") or ""),
|
|||
|
|
int(item["meta"].get("seq") or 0)))
|
|||
|
|
initial_hits = scored[:max(1, min(24, int(initial_top_k or 6)))]
|
|||
|
|
widened = max(int(initial_top_k or 6), int(widened_top_k or 24))
|
|||
|
|
# Mirror DSH Worker: widen only when the first recall page is sparse.
|
|||
|
|
hits = scored[:max(1, min(64, widened))] if len(initial_hits) < int(initial_top_k or 6) else initial_hits
|
|||
|
|
selected: list[dict] = []
|
|||
|
|
used_chars = 0
|
|||
|
|
for hit in hits:
|
|||
|
|
remaining = int(char_budget or 12000) - used_chars
|
|||
|
|
if remaining < 240 and selected:
|
|||
|
|
break
|
|||
|
|
excerpt = source_excerpt(hit.get("text", ""), query_text,
|
|||
|
|
limit=min(2400, max(240, remaining)))
|
|||
|
|
cost = len(excerpt)
|
|||
|
|
if selected and used_chars + cost > int(char_budget or 12000):
|
|||
|
|
continue
|
|||
|
|
selected.append(hit)
|
|||
|
|
used_chars += cost
|
|||
|
|
if used_chars >= int(char_budget or 12000):
|
|||
|
|
break
|
|||
|
|
reserved_ids = [str(hit["id"]) for hit in selected]
|
|||
|
|
for cid in reserved_ids:
|
|||
|
|
available.pop(cid, None)
|
|||
|
|
node_results.append({
|
|||
|
|
"order": index,
|
|||
|
|
"query": query_text[:500],
|
|||
|
|
"candidate_count": len(scored),
|
|||
|
|
"reserved_count": len(reserved_ids),
|
|||
|
|
"reserved_characters": used_chars,
|
|||
|
|
"chunk_ids": reserved_ids,
|
|||
|
|
"files": sorted({str(hit["meta"].get("file") or "") for hit in selected}),
|
|||
|
|
})
|
|||
|
|
shortfalls = [item for item in node_results if item["reserved_count"] <= 0]
|
|||
|
|
return {
|
|||
|
|
"ok": not shortfalls,
|
|||
|
|
"total_chunks": len(pool),
|
|||
|
|
"remaining_chunks": len(available),
|
|||
|
|
"nodes": node_results,
|
|||
|
|
"shortfalls": shortfalls,
|
|||
|
|
"retrieval_mode": "keyword-light",
|
|||
|
|
"vector_backend": "disabled",
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _lightweight_terms(query: str) -> list[str]:
|
|||
|
|
"""Extract conservative Chinese/English/numeric terms without a tokenizer.
|
|||
|
|
|
|||
|
|
Chinese contiguous phrases are kept in full and expanded to bigrams. This is
|
|||
|
|
intentionally simple and deterministic: it provides usable source-grounded
|
|||
|
|
recall in lightweight mode while avoiding embedding models, Chroma and torch.
|
|||
|
|
"""
|
|||
|
|
text = unicodedata.normalize("NFKC", str(query or "")).lower()
|
|||
|
|
terms: list[str] = []
|
|||
|
|
|
|||
|
|
def add(value: str) -> None:
|
|||
|
|
value = value.strip().lower()
|
|||
|
|
if len(value) >= 2 and value not in terms:
|
|||
|
|
terms.append(value)
|
|||
|
|
|
|||
|
|
for match in _LIGHTWEIGHT_NUMBER_RE.finditer(text):
|
|||
|
|
add(match.group(0))
|
|||
|
|
for match in _LIGHTWEIGHT_CJK_RE.finditer(text):
|
|||
|
|
token = match.group(0)
|
|||
|
|
add(token)
|
|||
|
|
if len(token) > 2:
|
|||
|
|
for idx in range(len(token) - 1):
|
|||
|
|
add(token[idx:idx + 2])
|
|||
|
|
for match in _LIGHTWEIGHT_ASCII_RE.finditer(text):
|
|||
|
|
add(match.group(0))
|
|||
|
|
return terms[:80]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _lightweight_score(text: str, terms: list[str]) -> tuple[float, int]:
|
|||
|
|
body = unicodedata.normalize("NFKC", text or "").lower()
|
|||
|
|
score = 0.0
|
|||
|
|
matched = 0
|
|||
|
|
for term in terms:
|
|||
|
|
count = body.count(term)
|
|||
|
|
if count <= 0:
|
|||
|
|
continue
|
|||
|
|
matched += 1
|
|||
|
|
if _LIGHTWEIGHT_CJK_RE.fullmatch(term):
|
|||
|
|
weight = 5.0 if len(term) >= 4 else (3.0 if len(term) == 3 else 1.0)
|
|||
|
|
elif _LIGHTWEIGHT_NUMBER_RE.fullmatch(term):
|
|||
|
|
weight = 5.0
|
|||
|
|
elif len(term) >= 5:
|
|||
|
|
weight = 3.0
|
|||
|
|
else:
|
|||
|
|
weight = 2.0
|
|||
|
|
score += min(count, 6) * weight
|
|||
|
|
if matched and terms:
|
|||
|
|
# Prefer chunks matching several independent query concepts, while term
|
|||
|
|
# frequency remains the dominant signal.
|
|||
|
|
score *= 1.0 + 0.35 * (matched / min(len(terms), 12))
|
|||
|
|
return score, matched
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _lightweight_keyword_search(session: Session, kb_ids: list[int], query: str, k: int = 6,
|
|||
|
|
excluded_chunk_ids: dict[int, set[str]] | None = None, *,
|
|||
|
|
campaign_id: int = 0, trace_id: str = "", user_id: int = 0,
|
|||
|
|
organization_id: int = 1, category_ids: list[int] | None = None,
|
|||
|
|
file_whitelists: dict[int, set[str]] | None = None,
|
|||
|
|
validate_access: bool = True) -> list[dict]:
|
|||
|
|
"""Keyword recall directly from registered source files when vectors are off.
|
|||
|
|
|
|||
|
|
The workbench still owns source files, KnowledgeFile registry, category
|
|||
|
|
filters, quarantine, per-Agent consumed chunks and citation reservation.
|
|||
|
|
Parsing/chunking is performed on demand and never persists local vectors.
|
|||
|
|
"""
|
|||
|
|
started = time.perf_counter()
|
|||
|
|
kids = list(dict.fromkeys(int(kid) for kid in kb_ids))
|
|||
|
|
query = str(query or "").strip()
|
|||
|
|
k = max(1, min(100, int(k or 6)))
|
|||
|
|
if not kids or not query:
|
|||
|
|
return []
|
|||
|
|
if len(query) > 10000:
|
|||
|
|
raise ValueError("检索问题不能超过 10000 字")
|
|||
|
|
terms = _lightweight_terms(query)
|
|||
|
|
if not terms:
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
campaign = session.get(Campaign, campaign_id) if campaign_id else None
|
|||
|
|
resolved_user = int(user_id or (campaign.owner_user_id if campaign else 0) or 0)
|
|||
|
|
resolved_org = int((campaign.organization_id if campaign and not user_id else organization_id) or 1)
|
|||
|
|
exclusions = {int(kid): set(values) for kid, values in (excluded_chunk_ids or {}).items()}
|
|||
|
|
specified_ranges = {int(kid): set(names) for kid, names in (file_whitelists or {}).items()}
|
|||
|
|
|
|||
|
|
actor = None
|
|||
|
|
if validate_access and resolved_user:
|
|||
|
|
from ..models import AppUser
|
|||
|
|
actor = session.exec(select(AppUser).where(AppUser.id == resolved_user)
|
|||
|
|
.execution_options(populate_existing=True)).first()
|
|||
|
|
if actor is None or not actor.enabled:
|
|||
|
|
raise ValueError("当前用户不存在或已停用，无法检索知识库")
|
|||
|
|
if int(actor.organization_id or 1) != resolved_org:
|
|||
|
|
raise ValueError("当前用户与知识库检索组织不一致")
|
|||
|
|
|
|||
|
|
candidates: list[dict] = []
|
|||
|
|
for kid in kids:
|
|||
|
|
kb = session.exec(select(KnowledgeBase).where(KnowledgeBase.id == kid)
|
|||
|
|
.execution_options(populate_existing=True)).first()
|
|||
|
|
if kb is None or int(kb.organization_id) not in {0, resolved_org}:
|
|||
|
|
raise ValueError("知识库不存在或当前用户没有访问权限")
|
|||
|
|
if kb.index_status == "indexing":
|
|||
|
|
raise ValueError("知识库正在导入，请完成后重新检索")
|
|||
|
|
if validate_access and actor:
|
|||
|
|
from ..access import can_access, canvas_allows_resource
|
|||
|
|
from ..models import KnowledgeBaseGrant
|
|||
|
|
grant = session.exec(select(KnowledgeBaseGrant).where(
|
|||
|
|
KnowledgeBaseGrant.kb_id == kid, KnowledgeBaseGrant.user_id == actor.id,
|
|||
|
|
KnowledgeBaseGrant.enabled == True)).first() # noqa: E712
|
|||
|
|
if not (can_access(kb, actor, "use", grant.permission if grant else None)
|
|||
|
|
or canvas_allows_resource(session, kb, actor, "use")):
|
|||
|
|
raise ValueError("知识库不存在或当前用户没有访问权限")
|
|||
|
|
|
|||
|
|
whitelist = _category_file_whitelist(session, kid, list(category_ids or []))
|
|||
|
|
if kid in specified_ranges:
|
|||
|
|
specified = specified_ranges[kid]
|
|||
|
|
whitelist = specified if whitelist is None else whitelist & specified
|
|||
|
|
# An explicitly empty set means no eligible file; it must not mean all.
|
|||
|
|
if whitelist is not None and not whitelist:
|
|||
|
|
continue
|
|||
|
|
quarantined = file_exclusions(session, kid)
|
|||
|
|
if whitelist is not None:
|
|||
|
|
whitelist = whitelist - quarantined
|
|||
|
|
if not whitelist:
|
|||
|
|
continue
|
|||
|
|
source = _source_dir(kb)
|
|||
|
|
records = session.exec(select(KnowledgeFile).where(KnowledgeFile.kb_id == kid)).all()
|
|||
|
|
for record in sorted(records, key=lambda item: item.file_name):
|
|||
|
|
file_name = str(record.file_name or "")
|
|||
|
|
if whitelist is not None and file_name not in whitelist:
|
|||
|
|
continue
|
|||
|
|
if file_name in quarantined:
|
|||
|
|
continue
|
|||
|
|
path = source / file_name
|
|||
|
|
if not path.is_file() or path.suffix.lower() not in _LIGHTWEIGHT_TEXT_SUFFIXES:
|
|||
|
|
continue
|
|||
|
|
try:
|
|||
|
|
units = parse.parse_file(path)
|
|||
|
|
chunks = chunk.chunk_texts(units, target=_lightweight_chunk_target(path, kb))
|
|||
|
|
except Exception as exc:
|
|||
|
|
logger.warning("轻量知识库解析失败，已跳过 %s/%s：%s", kid, file_name, exc)
|
|||
|
|
continue
|
|||
|
|
excluded = exclusions.get(kid, set())
|
|||
|
|
for seq, item in enumerate(chunks):
|
|||
|
|
text = str(item.get("text") or "").strip()
|
|||
|
|
if not text:
|
|||
|
|
continue
|
|||
|
|
meta = dict(item.get("meta") or {})
|
|||
|
|
heading = str(meta.get("heading") or "")
|
|||
|
|
chunk_id = _lightweight_chunk_id(
|
|||
|
|
kid, file_name, str(record.content_hash or ""), heading, seq, text)
|
|||
|
|
if chunk_id in excluded:
|
|||
|
|
continue
|
|||
|
|
score, matched = _lightweight_score(text, terms)
|
|||
|
|
if score <= 0 or matched <= 0:
|
|||
|
|
continue
|
|||
|
|
meta.update({"kb_id": kid, "file": file_name, "heading": heading,
|
|||
|
|
"seq": seq, "retrieval_mode": "keyword-light"})
|
|||
|
|
candidates.append({
|
|||
|
|
"id": chunk_id,
|
|||
|
|
"text": text,
|
|||
|
|
"meta": meta,
|
|||
|
|
"kb_id": kid,
|
|||
|
|
"distance": 0.0,
|
|||
|
|
"_score": round(score, 3),
|
|||
|
|
"_matched": matched,
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
candidates.sort(key=lambda item: (-float(item["_score"]), -int(item["_matched"]),
|
|||
|
|
int(item["kb_id"]),
|
|||
|
|
str(item["meta"].get("file") or ""),
|
|||
|
|
int(item["meta"].get("seq") or 0)))
|
|||
|
|
results = candidates[:k]
|
|||
|
|
elapsed_ms = round((time.perf_counter() - started) * 1000, 2)
|
|||
|
|
timings = {"validation_ms": 0.0, "embedding_ms": 0.0, "recall_ms": elapsed_ms,
|
|||
|
|
"postprocess_ms": 0.0, "total_ms": elapsed_ms, "retrieval_mode": "keyword-light",
|
|||
|
|
"vector_backend": "disabled", "candidate_count": len(candidates)}
|
|||
|
|
for item in results:
|
|||
|
|
item["retrieval_timings"] = timings
|
|||
|
|
item.pop("_score", None)
|
|||
|
|
item.pop("_matched", None)
|
|||
|
|
logger.info("轻量关键词检索完成：%s 个库，query=%r，%s 个候选，%s 条结果，%.3fs",
|
|||
|
|
len(kids), query[:80], len(candidates), len(results), elapsed_ms / 1000)
|
|||
|
|
return results
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _category_file_whitelist(session: Session, kb_id: int, category_ids: list[int]) -> set[str] | None:
|
|||
|
|
"""分类 → 允许的文件名集合；空 list 返回 None（不限制）。"""
|
|||
|
|
if not category_ids:
|
|||
|
|
return None
|
|||
|
|
ids = list(category_ids)
|
|||
|
|
# 含子孙分类
|
|||
|
|
cats = session.exec(select(KnowledgeCategory).where(KnowledgeCategory.kb_id == kb_id)).all()
|
|||
|
|
by_parent: dict = {}
|
|||
|
|
for c in cats:
|
|||
|
|
by_parent.setdefault(c.parent_id, []).append(c.id)
|
|||
|
|
expand = list(ids)
|
|||
|
|
queue = list(ids)
|
|||
|
|
while queue:
|
|||
|
|
cur = queue.pop()
|
|||
|
|
for child in by_parent.get(cur, []):
|
|||
|
|
if child not in expand:
|
|||
|
|
expand.append(child)
|
|||
|
|
queue.append(child)
|
|||
|
|
files = session.exec(select(KnowledgeFile).where(
|
|||
|
|
KnowledgeFile.kb_id == kb_id)).all()
|
|||
|
|
whitelist = set()
|
|||
|
|
for f in files:
|
|||
|
|
cat = f.category_id if f.category_id is not None else f.ai_category_id
|
|||
|
|
if cat in expand:
|
|||
|
|
whitelist.add(f.file_name)
|
|||
|
|
return whitelist
|
|||
|
|
|
|||
|
|
|
|||
|
|
def search_kbs(session: Session, kb_ids: list[int], query: str, k: int = 6,
|
|||
|
|
excluded_chunk_ids: dict[int, set[str]] | None = None, *,
|
|||
|
|
campaign_id: int = 0, trace_id: str = "", user_id: int = 0,
|
|||
|
|
organization_id: int = 1, category_ids: list[int] | None = None,
|
|||
|
|
file_whitelists: dict[int, set[str]] | None = None) -> list[dict]:
|
|||
|
|
"""独立中文关键词与向量召回、融合精排；缓存不绕过权限或消费过滤。"""
|
|||
|
|
global _LAST_RETRIEVAL_METRIC
|
|||
|
|
_LAST_RETRIEVAL_METRIC = None
|
|||
|
|
import logging
|
|||
|
|
import time
|
|||
|
|
from . import hybrid
|
|||
|
|
from ..access import can_access, canvas_allows_resource
|
|||
|
|
from ..models import AppUser, KnowledgeBaseGrant
|
|||
|
|
|
|||
|
|
started = time.perf_counter()
|
|||
|
|
validation_started = started
|
|||
|
|
if not kb_ids or not str(query or "").strip():
|
|||
|
|
return []
|
|||
|
|
if backend.is_disabled():
|
|||
|
|
return _lightweight_keyword_search(
|
|||
|
|
session, kb_ids, query, k=k, excluded_chunk_ids=excluded_chunk_ids,
|
|||
|
|
campaign_id=campaign_id, trace_id=trace_id, user_id=user_id,
|
|||
|
|
organization_id=organization_id, category_ids=category_ids,
|
|||
|
|
file_whitelists=file_whitelists, validate_access=True)
|
|||
|
|
k = max(1, min(100, int(k)))
|
|||
|
|
query = str(query).strip()
|
|||
|
|
if len(query) > 10000:
|
|||
|
|
raise ValueError("检索问题不能超过 10000 字")
|
|||
|
|
ids = list(dict.fromkeys(int(kid) for kid in kb_ids))
|
|||
|
|
campaign = session.get(Campaign, campaign_id) if campaign_id else None
|
|||
|
|
resolved_user = int(user_id or (campaign.owner_user_id if campaign else 0) or 0)
|
|||
|
|
resolved_org = int((campaign.organization_id if campaign and not user_id else organization_id) or 1)
|
|||
|
|
# A real requester is checked on every call, including cache hits. Internal
|
|||
|
|
# jobs without an account still keep the explicit organization boundary.
|
|||
|
|
actor = session.exec(select(AppUser).where(AppUser.id == resolved_user)
|
|||
|
|
.execution_options(populate_existing=True)).first() if resolved_user else None
|
|||
|
|
if resolved_user and (actor is None or not actor.enabled):
|
|||
|
|
raise ValueError("当前用户不存在或已停用，无法检索知识库")
|
|||
|
|
if actor and int(actor.organization_id or 1) != resolved_org:
|
|||
|
|
raise ValueError("当前用户与知识库检索组织不一致")
|
|||
|
|
kb_rows, allowed, scopes = {}, {}, []
|
|||
|
|
# Authorization is part of candidate-pool identity. Keep a compact
|
|||
|
|
# signature alongside scopes so a changed grant cannot reuse a prior pool.
|
|||
|
|
permission_signature = []
|
|||
|
|
# Validate before embedding and before looking in any query cache.
|
|||
|
|
for kid in ids:
|
|||
|
|
kb = session.exec(select(KnowledgeBase).where(KnowledgeBase.id == kid)
|
|||
|
|
.execution_options(populate_existing=True)).first()
|
|||
|
|
if kb is None or int(kb.organization_id) not in {0, resolved_org}:
|
|||
|
|
raise ValueError("知识库不存在或当前用户没有访问权限")
|
|||
|
|
if kb.index_status == "indexing":
|
|||
|
|
raise ValueError("知识库正在建立索引，请完成后重新检索")
|
|||
|
|
if actor:
|
|||
|
|
grant = session.exec(select(KnowledgeBaseGrant).where(
|
|||
|
|
KnowledgeBaseGrant.kb_id == kid, KnowledgeBaseGrant.user_id == actor.id,
|
|||
|
|
KnowledgeBaseGrant.enabled == True)).first() # noqa: E712
|
|||
|
|
if not (can_access(kb, actor, 'use', grant.permission if grant else None)
|
|||
|
|
or canvas_allows_resource(session, kb, actor, 'use')):
|
|||
|
|
raise ValueError("知识库不存在或当前用户没有访问权限")
|
|||
|
|
permission_signature.append((kid, str(getattr(grant, 'permission', '') or 'canvas')))
|
|||
|
|
else:
|
|||
|
|
permission_signature.append((kid, 'internal'))
|
|||
|
|
kb_rows[kid] = kb
|
|||
|
|
quarantined = file_exclusions(session, kid)
|
|||
|
|
whitelist = _category_file_whitelist(session, kid, list(category_ids or []))
|
|||
|
|
if file_whitelists is not None and kid in file_whitelists:
|
|||
|
|
specified = set(file_whitelists[kid])
|
|||
|
|
whitelist = specified if whitelist is None else whitelist & specified
|
|||
|
|
if quarantined:
|
|||
|
|
whitelist = (set(store.count_by_file(kid)) if whitelist is None else whitelist) - quarantined
|
|||
|
|
allowed[kid] = whitelist
|
|||
|
|
scopes.append((kid, hybrid.revision(kid), None if whitelist is None else sorted(whitelist)))
|
|||
|
|
validation_ms = (time.perf_counter() - validation_started) * 1000
|
|||
|
|
exclusions = {int(kid): set(values) for kid, values in (excluded_chunk_ids or {}).items()}
|
|||
|
|
# Bound work even when a long-running Agent has consumed many thousands of
|
|||
|
|
# chunks. The candidate pool itself is cached before consumption filtering.
|
|||
|
|
excluded_count = max((len(exclusions.get(kid, set())) for kid in ids), default=0)
|
|||
|
|
# Keep the first-stage pool small enough for fast disk/vector reads. A
|
|||
|
|
# bounded allowance covers recently consumed chunks without letting a
|
|||
|
|
# heavily used Agent turn every query into a 2k-row scan.
|
|||
|
|
needed = max(32, k * 4) + min(excluded_count, 256)
|
|||
|
|
# Buckets allow the same cached pool to survive a few newly consumed chunks.
|
|||
|
|
limit = min(512, 1 << (needed - 1).bit_length())
|
|||
|
|
model = str(getattr(embed, "EMBEDDING_MODEL", "local-embedding"))
|
|||
|
|
cache_model = model + ":" + str(getattr(embed, "EMBEDDING_POLICY", "default"))
|
|||
|
|
# Include lexical-index readiness in the cache namespace. A first query
|
|||
|
|
# may intentionally use vector-only recall while a large FTS index is being
|
|||
|
|
# built in the background; once that index becomes ready the key changes so
|
|||
|
|
# later requests immediately receive the hybrid (vector + keyword) result
|
|||
|
|
# instead of reusing the temporary vector-only pool for its full TTL.
|
|||
|
|
def _fts_scope_ready(kid, rev):
|
|||
|
|
files = allowed[kid]
|
|||
|
|
# A full-library index is also sufficient for a category scope: the
|
|||
|
|
# lexical channel filters its metadata after matching.
|
|||
|
|
return hybrid.index_available(resolved_org, kid, rev, files) or (
|
|||
|
|
files is not None and hybrid.index_available(resolved_org, kid, rev, None))
|
|||
|
|
fts_ready = all(_fts_scope_ready(kid, rev) for kid, rev, _ in scopes)
|
|||
|
|
cache_model += ":fts-ready" if fts_ready else ":fts-pending"
|
|||
|
|
retrieval_mode = "hybrid" if fts_ready else "vector"
|
|||
|
|
key = hybrid.candidate_key(
|
|||
|
|
resolved_org, resolved_user, scopes, query, limit, cache_model,
|
|||
|
|
mode=retrieval_mode, top_k=k, permissions=permission_signature,
|
|||
|
|
index_version=":".join(str(rev) for _, rev, _ in scopes),
|
|||
|
|
)
|
|||
|
|
pool = hybrid.cache_get(key)
|
|||
|
|
cache_hit = pool is not None
|
|||
|
|
embedding_ms = 0.0
|
|||
|
|
recall_ms = 0.0
|
|||
|
|
if pool is None:
|
|||
|
|
embedding_started = time.perf_counter()
|
|||
|
|
qv = embed.embed_query(query)
|
|||
|
|
embedding_ms = (time.perf_counter() - embedding_started) * 1000
|
|||
|
|
# Embed workers expose phase-level diagnostics (model load, tokenizer,
|
|||
|
|
# inference and serialization). Keep these alongside retrieval timing
|
|||
|
|
# without changing the vector API or test doubles.
|
|||
|
|
embedding_metric = getattr(embed, "last_embedding_metric", lambda: {})()
|
|||
|
|
_record_query_embedding_usage(
|
|||
|
|
idempotency_key=f"embedding:query:{trace_id or uuid.uuid4().hex}:{hashlib.sha256(query.encode('utf-8')).hexdigest()}",
|
|||
|
|
user_id=resolved_user, organization_id=resolved_org, campaign_id=campaign_id,
|
|||
|
|
trace_id=trace_id, model=model, quantity=estimate_tokens(query),
|
|||
|
|
metadata={"kb_ids": ids, "operation": "query"})
|
|||
|
|
def recall_scope(scope):
|
|||
|
|
"""召回单个知识库；独立 I/O 并行，避免多库检索按库串行等待。"""
|
|||
|
|
kid, rev, _ = scope
|
|||
|
|
whitelist = allowed[kid]
|
|||
|
|
try:
|
|||
|
|
scope_started = time.perf_counter()
|
|||
|
|
# Apply the same metadata scope to BOTH independent recalls.
|
|||
|
|
vhits = store.search(kid, qv, k=limit, **({'file_names': whitelist} if whitelist is not None else {}))
|
|||
|
|
vector_metric = getattr(store, "last_search_metric", lambda: {})()
|
|||
|
|
# lexical_search itself is bounded: it reuses the ready
|
|||
|
|
# full-library index for category scopes, or returns vector-only
|
|||
|
|
# while an absent index is being prewarmed in the background.
|
|||
|
|
lhits = hybrid.lexical_search(resolved_org, kid, rev, query, limit, whitelist)
|
|||
|
|
lexical_metric = getattr(hybrid, "last_lexical_metric", lambda: {})()
|
|||
|
|
channels = []
|
|||
|
|
for hits in (vhits, lhits):
|
|||
|
|
rows = []
|
|||
|
|
for channel_rank, item in enumerate(hits, 1):
|
|||
|
|
if whitelist is not None and (item.get('meta') or {}).get('file') not in whitelist:
|
|||
|
|
raise RuntimeError('检索返回了所选分类以外的文件')
|
|||
|
|
item = dict(item)
|
|||
|
|
item['kb_id'], item['kb_name'] = kid, kb_rows[kid].name
|
|||
|
|
item['channel_rank'] = channel_rank
|
|||
|
|
rows.append(item)
|
|||
|
|
channels.append(rows)
|
|||
|
|
if hybrid.revision(kid) != rev:
|
|||
|
|
raise RuntimeError('知识库索引正在更新，请稍后重新检索')
|
|||
|
|
return channels, {
|
|||
|
|
"kb_id": kid,
|
|||
|
|
"kb_name": kb_rows[kid].name,
|
|||
|
|
"collection_get_ms": vector_metric.get("collection_get_ms", 0.0),
|
|||
|
|
"metadata_filter_ms": round(float(vector_metric.get("metadata_filter_ms", 0.0) or 0.0)
|
|||
|
|
+ float(lexical_metric.get("metadata_filter_ms", 0.0) or 0.0), 2),
|
|||
|
|
"vector_ms": vector_metric.get("vector_ms", 0.0),
|
|||
|
|
"fts_ms": lexical_metric.get("fts_ms", 0.0),
|
|||
|
|
"fts_connection_reused": bool(lexical_metric.get("connection_reused", False)),
|
|||
|
|
"scope_ms": round((time.perf_counter() - scope_started) * 1000, 2),
|
|||
|
|
"warm_state": lexical_metric.get("warm_state") or hybrid.warm_state(resolved_org, kid, rev, whitelist),
|
|||
|
|
}
|
|||
|
|
except Exception as exc:
|
|||
|
|
logging.getLogger('yawen.rag.search').exception('知识库 %s 检索失败', kid)
|
|||
|
|
raise RuntimeError(f'知识库 {kid} 检索失败：{exc}') from exc
|
|||
|
|
|
|||
|
|
# Vector and FTS calls are network/disk bound. Running the independent
|
|||
|
|
# knowledge-base scopes concurrently cuts multi-library latency to the
|
|||
|
|
# slowest scope while keeping the same bounded candidate limit.
|
|||
|
|
recall_started = time.perf_counter()
|
|||
|
|
if len(scopes) > 1:
|
|||
|
|
# The process-wide bounded pool keeps recall concurrency stable and
|
|||
|
|
# lets each worker reuse its read-only FTS connection across
|
|||
|
|
# requests. ``executor.map`` still preserves scope ordering.
|
|||
|
|
recalled = list(_recall_executor.map(recall_scope, scopes))
|
|||
|
|
elif scopes:
|
|||
|
|
recalled = [recall_scope(scopes[0])]
|
|||
|
|
else:
|
|||
|
|
recalled = []
|
|||
|
|
scope_metrics = [item[1] for item in recalled]
|
|||
|
|
recalled_channels = [item[0] for item in recalled]
|
|||
|
|
vector = [item for channels in recalled_channels for item in channels[0]]
|
|||
|
|
lexical = [item for channels in recalled_channels for item in channels[1]]
|
|||
|
|
recall_ms = (time.perf_counter() - recall_started) * 1000
|
|||
|
|
# Channel scores are ordered across libraries before fusion; incomparable
|
|||
|
|
# per-library FTS BM25 scores use reciprocal positions instead of raw sums.
|
|||
|
|
vector.sort(key=lambda row: (float(row.get('distance') or 0.0), int(row['kb_id']), str(row['id'])))
|
|||
|
|
lexical.sort(key=lambda row: (row['channel_rank'], int(row['kb_id']), str(row['id'])))
|
|||
|
|
pool = {'vector': vector, 'lexical': lexical}
|
|||
|
|
hybrid.cache_put(key, pool)
|
|||
|
|
else:
|
|||
|
|
# Preserve per-library visibility on cache hits. Stage durations are
|
|||
|
|
# zero because no vector/FTS I/O ran, while warm state still tells the
|
|||
|
|
# operator whether the cached pool was backed by a ready index.
|
|||
|
|
scope_metrics = [{
|
|||
|
|
"kb_id": kid, "kb_name": kb_rows[kid].name,
|
|||
|
|
"collection_get_ms": 0.0, "metadata_filter_ms": 0.0,
|
|||
|
|
"vector_ms": 0.0, "fts_ms": 0.0, "scope_ms": 0.0,
|
|||
|
|
"fts_connection_reused": False,
|
|||
|
|
"warm_state": hybrid.warm_state(resolved_org, kid, rev, allowed[kid]),
|
|||
|
|
} for kid, rev, _ in scopes]
|
|||
|
|
postprocess_started = time.perf_counter()
|
|||
|
|
# Re-check consumption BEFORE bounded candidate reranking on every request.
|
|||
|
|
# Cached pools are never final, previously authorized answers.
|
|||
|
|
rerank_limit = min(128, max(32, k * 2))
|
|||
|
|
filtered = {}
|
|||
|
|
for channel in ('vector', 'lexical'):
|
|||
|
|
filtered[channel] = [row for row in pool[channel]
|
|||
|
|
if row['id'] not in exclusions.get(int(row['kb_id']), set())][:rerank_limit]
|
|||
|
|
candidates = hybrid.fuse(filtered['vector'], filtered['lexical'], query)
|
|||
|
|
# Optional local Cross-Encoder rerank. The module is deliberately lazy and
|
|||
|
|
# returns the fused order unchanged when no model is configured or loading
|
|||
|
|
# fails, so the default offline deployment keeps the existing latency and
|
|||
|
|
# behaviour guarantees.
|
|||
|
|
from . import rerank as _rerank
|
|||
|
|
candidates = _rerank.rerank(candidates, query, limit=min(64, max(16, k * 4)))
|
|||
|
|
rerank_metric = _rerank.last_metric()
|
|||
|
|
fusion_metric = hybrid.last_fusion_metric() if not cache_hit else {}
|
|||
|
|
for kid, rev, _ in scopes:
|
|||
|
|
if hybrid.revision(kid) != rev:
|
|||
|
|
raise RuntimeError('知识库索引正在更新，请稍后重新检索')
|
|||
|
|
results = hybrid.finalize(candidates, k, exclusions, allowed)
|
|||
|
|
postprocess_ms = (time.perf_counter() - postprocess_started) * 1000
|
|||
|
|
timings = {
|
|||
|
|
"cache_hit": cache_hit,
|
|||
|
|
"validation_ms": round(validation_ms, 2),
|
|||
|
|
"embedding_ms": round(embedding_ms, 2),
|
|||
|
|
"recall_ms": round(recall_ms, 2),
|
|||
|
|
"postprocess_ms": round(postprocess_ms, 2),
|
|||
|
|
"total_ms": round((time.perf_counter() - started) * 1000, 2),
|
|||
|
|
"retrieval_scopes": scope_metrics,
|
|||
|
|
"warm_state": {str(row.get("kb_id")): row.get("warm_state") for row in scope_metrics},
|
|||
|
|
"merge_ms": fusion_metric.get("merge_ms", 0.0),
|
|||
|
|
"dedup_ms": fusion_metric.get("dedup_ms", 0.0),
|
|||
|
|
"rrf_ms": fusion_metric.get("rrf_ms", 0.0),
|
|||
|
|
"rerank_enabled": bool(rerank_metric.get("enabled", False)),
|
|||
|
|
"rerank_model": rerank_metric.get("model", ""),
|
|||
|
|
"rerank_model_loaded": bool(rerank_metric.get("model_loaded", False)),
|
|||
|
|
"rerank_candidate_count": int(rerank_metric.get("candidate_count", 0) or 0),
|
|||
|
|
"rerank_count": int(rerank_metric.get("reranked_count", 0) or 0),
|
|||
|
|
"rerank_model_load_ms": rerank_metric.get("model_load_ms", 0.0),
|
|||
|
|
"rerank_inference_ms": rerank_metric.get("inference_ms", 0.0),
|
|||
|
|
"rerank_fallback": bool(rerank_metric.get("fallback", True)),
|
|||
|
|
}
|
|||
|
|
if pool is not None:
|
|||
|
|
metric = locals().get("embedding_metric") or {}
|
|||
|
|
if metric:
|
|||
|
|
query_cache = getattr(embed, "query_cache_snapshot", lambda: {})()
|
|||
|
|
timings.update({
|
|||
|
|
"embedding_source": metric.get("embedding_source"),
|
|||
|
|
"embedding_count": metric.get("embedding_count", 1),
|
|||
|
|
"embedding_model": metric.get("model", metric.get("model_name")),
|
|||
|
|
"embedding_version": metric.get("model_version", metric.get("version")),
|
|||
|
|
"embedding_pid": metric.get("pid", metric.get("process_id")),
|
|||
|
|
"embedding_dim": metric.get("embedding_dimension", metric.get("embedding_dim", 0)),
|
|||
|
|
"embedding_cache_hit": metric.get("cache_hit", False),
|
|||
|
|
"embedding_model_load_ms": metric.get("model_load_ms", 0.0),
|
|||
|
|
"embedding_tokenizer_ms": metric.get("tokenizer_ms", 0.0),
|
|||
|
|
"embedding_encode_ms": metric.get("encode_ms", metric.get("encode_inference_ms", 0.0)),
|
|||
|
|
"embedding_postprocess_ms": metric.get("postprocess_ms", metric.get("postprocess_serialization_ms", 0.0)),
|
|||
|
|
"embedding_process_start_ms": metric.get("process_start_ms", 0.0),
|
|||
|
|
"embedding_ipc_ms": metric.get("ipc_ms", 0.0),
|
|||
|
|
"embedding_total_ms": metric.get("total_ms", embedding_ms),
|
|||
|
|
"embedding_cache_size": query_cache.get("size", 0),
|
|||
|
|
"embedding_cache_hits": query_cache.get("hits", 0),
|
|||
|
|
"embedding_cache_misses": query_cache.get("misses", 0),
|
|||
|
|
"embedding_cache_hit_rate": query_cache.get("hit_rate", 0.0),
|
|||
|
|
})
|
|||
|
|
_record_retrieval_metrics(timings, len(results))
|
|||
|
|
for item in results:
|
|||
|
|
item["retrieval_timings"] = timings
|
|||
|
|
logging.getLogger('yawen.rag.search').info(
|
|||
|
|
'知识库检索完成：%s 个库，%s 个候选，%s 条结果，缓存%s，%.3f 秒',
|
|||
|
|
len(ids), len(candidates), len(results), '命中' if cache_hit else '未命中', time.perf_counter() - started)
|
|||
|
|
return results
|
|||
|
|
|
|||
|
|
|
|||
|
|
def import_folder(session: Session, kb_id: int, folder_zip: str, rel_base: str = "", progress_cb=None) -> dict:
|
|||
|
|
"""导入已解压的文件夹（含多层子文件夹）。
|
|||
|
|
file_paths 里每个条目形如 "相对路径\t绝对路径"，相对路径用于来源标注。"""
|
|||
|
|
entries = []
|
|||
|
|
for line in str(folder_zip).split("\n"):
|
|||
|
|
line = line.strip()
|
|||
|
|
if not line or "\t" not in line:
|
|||
|
|
continue
|
|||
|
|
rel, abs_path = line.rsplit("\t", 1)
|
|||
|
|
entries.append((rel, abs_path))
|
|||
|
|
if not entries:
|
|||
|
|
raise ValueError("文件夹里没有可导入的文件")
|
|||
|
|
|
|||
|
|
def cb(name: str, n: int):
|
|||
|
|
if progress_cb:
|
|||
|
|
progress_cb(name, n)
|
|||
|
|
|
|||
|
|
return import_kb_files(session, kb_id, [abs_path for _, abs_path in entries], progress_cb=cb)
|