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