"""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)