feat(server): add workbench API, DSH planner worker and RAG pipeline
This commit is contained in:
yi
@@ -0,0 +1,78 @@
|
||||
"""Vector backend configuration.
|
||||
|
||||
Default is ``disabled`` for lightweight local startup. Install optional local
|
||||
RAG dependencies and set RAG_VECTOR_BACKEND=chroma to restore local vectors.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from ..config import DATA_DIR
|
||||
|
||||
|
||||
def settings():
|
||||
path = DATA_DIR / 'rag-backend.json'
|
||||
try:
|
||||
config = json.loads(path.read_text(encoding='utf-8'))
|
||||
if not isinstance(config, dict):
|
||||
raise ValueError('知识索引配置必须是对象')
|
||||
except FileNotFoundError:
|
||||
config = {}
|
||||
value = os.environ.get('RAG_VECTOR_BACKEND', config.get('backend', 'disabled'))
|
||||
backend_name = str(value or 'disabled').strip().lower()
|
||||
if backend_name not in {'disabled', 'chroma', 'qdrant'}:
|
||||
raise ValueError('知识索引后端配置无效')
|
||||
return {**config, 'backend': backend_name}
|
||||
|
||||
|
||||
def configured_backend() -> str:
|
||||
return settings()['backend']
|
||||
|
||||
|
||||
def local_chroma_available() -> bool:
|
||||
return importlib.util.find_spec('chromadb') is not None
|
||||
|
||||
|
||||
def effective_backend() -> str:
|
||||
"""Return the backend that can actually run in this installation."""
|
||||
value = configured_backend()
|
||||
if value == 'chroma' and not local_chroma_available():
|
||||
return 'disabled'
|
||||
return value
|
||||
|
||||
|
||||
def is_disabled() -> bool:
|
||||
return effective_backend() == 'disabled'
|
||||
|
||||
|
||||
def is_qdrant():
|
||||
return effective_backend() == 'qdrant'
|
||||
|
||||
|
||||
def qdrant_settings():
|
||||
config = settings()
|
||||
url = str(os.environ.get('QDRANT_URL', config.get('url', ''))).rstrip('/')
|
||||
if not url.startswith(('http://', 'https://')):
|
||||
raise ValueError('请配置 Qdrant 服务地址')
|
||||
key = os.environ.get('QDRANT_API_KEY', '')
|
||||
if not key and config.get('api_key_file'):
|
||||
key = Path(config['api_key_file']).read_text(encoding='utf-8').strip()
|
||||
if not key:
|
||||
raise ValueError('请配置 Qdrant API 密钥')
|
||||
prefix = str(config.get('collection_prefix', 'yawan_'))
|
||||
if not prefix or not all(c.isalnum() or c in '_-' for c in prefix):
|
||||
raise ValueError('Qdrant 知识库命名空间无效')
|
||||
return url, key, prefix
|
||||
|
||||
|
||||
def namespace():
|
||||
value = effective_backend()
|
||||
if value == 'qdrant':
|
||||
config = settings()
|
||||
return json.dumps(['qdrant', str(os.environ.get('QDRANT_URL', config.get('url', ''))).rstrip('/'),
|
||||
str(config.get('collection_prefix', 'yawan_'))], separators=(',', ':'))
|
||||
if value == 'disabled':
|
||||
return 'disabled'
|
||||
return 'chroma' if local_chroma_available() else 'chroma-missing'
|
||||
@@ -0,0 +1,62 @@
|
||||
"""RAG 管道：标题感知的连续窗口分块。"""
|
||||
import re
|
||||
|
||||
TARGET = 500
|
||||
MIN_CHUNK = 120
|
||||
OVERLAP_RATIO = 0.15
|
||||
|
||||
def chunk_texts(units: list[dict], target: int = TARGET) -> list[dict]:
|
||||
"""连续合并相同元数据的相邻单元，避免字幕类文件被切成大量细碎向量。"""
|
||||
target = max(MIN_CHUNK, int(target or TARGET))
|
||||
chunks = []
|
||||
buf: list[str] = []
|
||||
buf_len = 0
|
||||
active_meta: dict | None = None
|
||||
|
||||
def flush() -> None:
|
||||
nonlocal buf, buf_len
|
||||
if buf:
|
||||
text = "".join(buf)
|
||||
if len(text) >= 30:
|
||||
chunks.append({"text": text, "meta": dict(active_meta or {})})
|
||||
buf, buf_len = [], 0
|
||||
|
||||
for u in units:
|
||||
text = str(u.get("text") or "")
|
||||
meta = dict(u.get("meta") or {})
|
||||
if not text:
|
||||
continue
|
||||
if active_meta is not None and meta != active_meta:
|
||||
flush()
|
||||
if active_meta is None or meta != active_meta:
|
||||
active_meta = meta
|
||||
# 按句子/标点切候选段
|
||||
parts = _split_sentences(text)
|
||||
for p in parts:
|
||||
if buf_len + len(p) > target * 1.15 and buf:
|
||||
chunks.append({"text": "".join(buf), "meta": dict(active_meta)})
|
||||
# 重叠：保留最后 15% 字符
|
||||
overlap = "".join(buf)
|
||||
cut = int(len(overlap) * OVERLAP_RATIO)
|
||||
keep = overlap[-cut:] if cut > 30 else ""
|
||||
buf = [keep] if keep else []
|
||||
buf_len = len(keep)
|
||||
buf.append(p)
|
||||
buf_len += len(p)
|
||||
flush()
|
||||
return chunks
|
||||
|
||||
def _split_sentences(text: str) -> list[str]:
|
||||
# 在句号/问号/感叹号/换行后切分，保持小段不小于 30 字
|
||||
parts = re.split(r"(?<=[。！？!?；;])\s*|\n+", text)
|
||||
merged, cur = [], ""
|
||||
for p in parts:
|
||||
if not p:
|
||||
continue
|
||||
cur += p
|
||||
if len(cur) >= 30:
|
||||
merged.append(cur)
|
||||
cur = ""
|
||||
if cur:
|
||||
merged.append(cur)
|
||||
return merged or [text]
|
||||
@@ -0,0 +1,41 @@
|
||||
"""Bounded source excerpts with relevant middle/tail evidence and honest omissions."""
|
||||
import re
|
||||
|
||||
|
||||
def source_excerpt(text: str, query: str = "", limit: int = 2400) -> str:
|
||||
text = str(text or "")
|
||||
if len(text) <= limit:
|
||||
return text
|
||||
marker = "\n[…此处省略，未展示内容不能作为已知事实…]\n"
|
||||
# Keep boundary conditions and conclusion as well as query-related evidence.
|
||||
edge = min(350, limit // 5)
|
||||
spans = [(0, edge), (len(text) - edge, len(text))]
|
||||
terms = re.findall(r"[A-Za-z0-9_]+|[\u3400-\u9fff]{2,}", str(query).lower())
|
||||
terms += [term[i:i + 2] for term in list(terms) if len(term) > 2
|
||||
for i in range(len(term) - 1)]
|
||||
lower = text.lower()
|
||||
candidates = []
|
||||
for term in dict.fromkeys(terms):
|
||||
position = lower.find(term)
|
||||
if position >= edge and position < len(text) - edge:
|
||||
candidates.append((position, len(term)))
|
||||
candidates.sort(key=lambda pair: (-pair[1], pair[0]))
|
||||
remaining = max(0, limit - 2 * edge - 3 * len(marker))
|
||||
for position, _ in candidates:
|
||||
if remaining < 120:
|
||||
break
|
||||
if any(start <= position < end for start, end in spans):
|
||||
continue
|
||||
size = min(remaining, 900)
|
||||
start = max(0, position - size // 3)
|
||||
spans.append((start, min(len(text), start + size)))
|
||||
remaining -= size + len(marker)
|
||||
if len(spans) == 2 and remaining:
|
||||
spans.append((edge, min(len(text) - edge, edge + remaining)))
|
||||
merged = []
|
||||
for start, end in sorted(spans):
|
||||
if merged and start <= merged[-1][1]:
|
||||
merged[-1] = (merged[-1][0], max(end, merged[-1][1]))
|
||||
else:
|
||||
merged.append((start, end))
|
||||
return marker.join(text[start:end] for start, end in merged)
|
||||
@@ -0,0 +1,110 @@
|
||||
"""知识库导入去重：只比较精确正文，不以语义相似度删除资料。"""
|
||||
from contextlib import contextmanager
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import threading
|
||||
import uuid
|
||||
|
||||
from ..config import DATA_DIR
|
||||
|
||||
DEDUP_DIR = DATA_DIR / "rag_dedup"
|
||||
_locks: dict[int, threading.RLock] = {}
|
||||
_locks_guard = threading.Lock()
|
||||
|
||||
|
||||
def file_hash(path: Path) -> str:
|
||||
"""流式原始文件指纹；大文件不会为计算哈希额外复制到内存。"""
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for block in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def normalized_text(value: str) -> str:
|
||||
# 保留字词、数字、标点及它们的顺序，仅统一排版空白。
|
||||
return " ".join(str(value or "").split())
|
||||
|
||||
|
||||
def document_hash(units: list[dict]) -> str:
|
||||
digest = hashlib.sha256()
|
||||
for unit in units:
|
||||
text = normalized_text(unit.get("text", ""))
|
||||
if not text:
|
||||
continue
|
||||
# 标题是正文语义的一部分；页码/字幕编号不是。
|
||||
heading = normalized_text((unit.get("meta") or {}).get("heading", ""))
|
||||
for value in (heading, text):
|
||||
encoded = value.encode("utf-8")
|
||||
digest.update(len(encoded).to_bytes(8, "big"))
|
||||
digest.update(encoded)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def unique_chunks(chunks: list[dict]) -> tuple[list[dict], int]:
|
||||
"""同一文件精确重复段只嵌入一次，保留首条来源与原始顺序。"""
|
||||
unique, seen = [], {}
|
||||
for ordinal, item in enumerate(chunks):
|
||||
key = (normalized_text((item.get("meta") or {}).get("heading", "")),
|
||||
normalized_text(item.get("text", "")))
|
||||
if key in seen:
|
||||
seen[key]["meta"]["duplicate_count"] += 1
|
||||
continue
|
||||
retained = {**item, "meta": dict(item.get("meta") or {})}
|
||||
retained["meta"].update(seq=ordinal, source_seq=ordinal, duplicate_count=0)
|
||||
seen[key] = retained
|
||||
unique.append(retained)
|
||||
return unique, len(chunks) - len(unique)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def library_lock(kb_id: int):
|
||||
"""线程与进程都串行写同一库，避免同时上传相同文件穿过去重。"""
|
||||
with _locks_guard:
|
||||
lock = _locks.setdefault(int(kb_id), threading.RLock())
|
||||
with lock:
|
||||
DEDUP_DIR.mkdir(parents=True, exist_ok=True)
|
||||
with (DEDUP_DIR / f"kb_{int(kb_id)}.lock").open("a") as handle:
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError: # Windows 保留线程保护；生产为 macOS/Linux。
|
||||
fcntl = None
|
||||
if fcntl:
|
||||
fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if fcntl:
|
||||
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
|
||||
|
||||
|
||||
class FingerprintCache:
|
||||
"""仅缓存解析指纹；当前原文件 SHA256 和现有索引仍独立校验。"""
|
||||
def __init__(self, kb_id: int):
|
||||
self.path = DEDUP_DIR / f"kb_{int(kb_id)}.json"
|
||||
try:
|
||||
data = json.loads(self.path.read_text(encoding="utf-8"))
|
||||
self.values = data.get("values", {}) if data.get("version") == 1 else {}
|
||||
except (OSError, ValueError, TypeError, AttributeError):
|
||||
self.values = {}
|
||||
|
||||
def get(self, raw_hash: str) -> str:
|
||||
value = self.values.get(raw_hash)
|
||||
return value if isinstance(value, str) and len(value) == 64 else ""
|
||||
|
||||
def set(self, raw_hash: str, body_hash: str):
|
||||
self.values[raw_hash] = body_hash
|
||||
|
||||
def save(self):
|
||||
temp = self.path.with_suffix(f".{uuid.uuid4().hex}.tmp")
|
||||
try:
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temp.write_text(json.dumps({"version": 1, "values": self.values}), encoding="utf-8")
|
||||
os.replace(temp, self.path)
|
||||
except OSError:
|
||||
# 只是加速缓存，不能让已成功的入库变为失败。
|
||||
pass
|
||||
finally:
|
||||
temp.unlink(missing_ok=True)
|
||||
@@ -0,0 +1,288 @@
|
||||
"""RAG 管道：Embedding（本地模型，ModelScope 下载，离线可用）"""
|
||||
import sys
|
||||
import threading
|
||||
import math
|
||||
import logging
|
||||
import time
|
||||
import os
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from ..config import MODELS_DIR, EMBEDDING_MODEL
|
||||
|
||||
_lock = threading.Lock()
|
||||
_encoder = None
|
||||
_encoder_name = None
|
||||
_encoder_load_ms = 0.0
|
||||
_metric_lock = threading.Lock()
|
||||
_last_embedding_metric: dict = {}
|
||||
_embedding_dimension = ""
|
||||
logger = logging.getLogger("yawen.rag.embed")
|
||||
|
||||
# This is deliberately a separate value from the model name. Operators can
|
||||
# bump it when replacing files in the model directory without changing the
|
||||
# configured model identifier (which also invalidates query-vector caches).
|
||||
EMBEDDING_MODEL_VERSION = os.environ.get("EMBEDDING_MODEL_VERSION", "1")
|
||||
|
||||
|
||||
def _set_metric(row: dict) -> None:
|
||||
"""Keep the most recent bounded embedding sample for the RAG envelope.
|
||||
|
||||
The RAG worker is single threaded today, but protecting this assignment
|
||||
makes the helper safe when a worker is configured with concurrent requests.
|
||||
No source text is retained in telemetry.
|
||||
"""
|
||||
global _last_embedding_metric
|
||||
clean = {
|
||||
"process_start_ms": 0.0,
|
||||
"ipc_ms": 0.0,
|
||||
"model_load_ms": 0.0,
|
||||
"tokenizer_ms": 0.0,
|
||||
"encode_ms": 0.0,
|
||||
"postprocess_ms": 0.0,
|
||||
"embedding_source": "sentence-transformers",
|
||||
"embedding_count": 0,
|
||||
"model": str(EMBEDDING_MODEL),
|
||||
"model_version": str(EMBEDDING_MODEL_VERSION),
|
||||
"pid": os.getpid(),
|
||||
"cache_hit": False,
|
||||
}
|
||||
clean.update({key: value for key, value in row.items() if value is not None})
|
||||
for key in ("process_start_ms", "ipc_ms", "model_load_ms", "tokenizer_ms",
|
||||
"encode_ms", "postprocess_ms"):
|
||||
try:
|
||||
clean[key] = round(float(clean[key]), 2)
|
||||
except (TypeError, ValueError):
|
||||
clean[key] = 0.0
|
||||
clean["embedding_count"] = int(clean.get("embedding_count") or 0)
|
||||
clean["pid"] = int(clean.get("pid") or os.getpid())
|
||||
with _metric_lock:
|
||||
_last_embedding_metric = clean
|
||||
|
||||
|
||||
def last_embedding_metric() -> dict:
|
||||
"""Return phase timings for the latest encode without exposing input text."""
|
||||
with _metric_lock:
|
||||
return dict(_last_embedding_metric)
|
||||
|
||||
|
||||
def query_cache_snapshot() -> dict:
|
||||
"""Expose the independent query-vector LRU counters for diagnostics."""
|
||||
info = _cached_query.cache_info()
|
||||
total = info.hits + info.misses
|
||||
return {
|
||||
"size": info.currsize, "max_size": info.maxsize,
|
||||
"hits": info.hits, "misses": info.misses,
|
||||
"hit_rate": round(info.hits / total, 4) if total else 0.0,
|
||||
"model": str(EMBEDDING_MODEL), "model_version": str(EMBEDDING_MODEL_VERSION),
|
||||
"dimension": int(_embedding_dimension or 0),
|
||||
"pid": os.getpid(),
|
||||
}
|
||||
|
||||
|
||||
def clear_query_cache() -> None:
|
||||
"""Clear only query vectors; retrieval candidate pools remain untouched."""
|
||||
_cached_query.cache_clear()
|
||||
|
||||
|
||||
def set_transport_timings(*, process_start_ms: float = 0.0, ipc_ms: float = 0.0) -> None:
|
||||
"""Attach process/IPC timings measured by the parent worker transport.
|
||||
|
||||
``embed.py`` cannot observe subprocess creation or pipe transfer itself;
|
||||
the worker pool can call this after a request and the next envelope then
|
||||
contains the complete set of phases. Keeping this optional preserves the
|
||||
existing embedding API for in-process callers and tests.
|
||||
"""
|
||||
metric = last_embedding_metric()
|
||||
metric.update(process_start_ms=process_start_ms, ipc_ms=ipc_ms)
|
||||
_set_metric(metric)
|
||||
|
||||
|
||||
def ensure_model_downloaded(model_name: str) -> Path:
|
||||
"""模型目录存在即用；否则从 ModelScope 下载（HF 被墙）"""
|
||||
safe = model_name.replace("/", "--")
|
||||
local = MODELS_DIR / safe
|
||||
if (local / "config.json").exists():
|
||||
return local
|
||||
try:
|
||||
from modelscope import snapshot_download
|
||||
print(f"[embed] 从 ModelScope 下载 {model_name} ...", file=sys.stderr)
|
||||
snapshot_download(model_name, local_dir=str(local))
|
||||
return local
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"模型下载失败（ModelScope）: {e}")
|
||||
|
||||
|
||||
def get_encoder():
|
||||
"""单例加载 SentenceTransformer（本地路径）"""
|
||||
global _encoder, _encoder_name, _encoder_load_ms
|
||||
if _encoder is not None and _encoder_name == EMBEDDING_MODEL:
|
||||
return _encoder
|
||||
with _lock:
|
||||
if _encoder is not None and _encoder_name == EMBEDDING_MODEL:
|
||||
return _encoder
|
||||
started = time.perf_counter()
|
||||
from sentence_transformers import SentenceTransformer
|
||||
local = ensure_model_downloaded(EMBEDDING_MODEL)
|
||||
print(f"[embed] 加载 {local} ...", file=sys.stderr)
|
||||
_encoder = SentenceTransformer(str(local))
|
||||
_encoder_name = EMBEDDING_MODEL
|
||||
_encoder_load_ms = (time.perf_counter() - started) * 1000
|
||||
return _encoder
|
||||
|
||||
|
||||
def prewarm() -> bool:
|
||||
"""Load the embedding model and run one tiny encode as a best-effort warmup.
|
||||
|
||||
Model loading is intentionally kept out of the API startup path. The
|
||||
caller should run this function from a daemon thread so a slow model
|
||||
download/load never delays health checks or request handling. Returning a
|
||||
boolean keeps this useful for startup diagnostics without making warmup
|
||||
failures fatal to the service.
|
||||
"""
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
encoder = get_encoder()
|
||||
# A real encode initializes tokenizer/runtime kernels as well as the
|
||||
# SentenceTransformer object. The result is discarded; query vectors
|
||||
# continue to use the normal cached path.
|
||||
encoder.encode(["系统启动预热"], normalize_embeddings=True,
|
||||
show_progress_bar=False, batch_size=1)
|
||||
logger.info("Embedding 模型预热完成：%s，耗时 %.3f 秒",
|
||||
EMBEDDING_MODEL, time.perf_counter() - started)
|
||||
return True
|
||||
except Exception:
|
||||
logger.exception("Embedding 模型预热失败：%s，服务将按需加载", EMBEDDING_MODEL)
|
||||
return False
|
||||
|
||||
|
||||
EMBEDDING_POLICY = "full-token-windows-v1"
|
||||
|
||||
|
||||
def text_windows(text: str, encoder=None) -> list[str]:
|
||||
"""Cover every source character without sending over-limit windows to BGE."""
|
||||
enc = encoder or get_encoder()
|
||||
tokenizer = enc.tokenizer
|
||||
maximum = int(enc.max_seq_length)
|
||||
def fits(value):
|
||||
return len(tokenizer.encode(value, add_special_tokens=True, truncation=False)) <= maximum
|
||||
if fits(text):
|
||||
return [text]
|
||||
# Character boundaries preserve exact original text; a binary search uses the
|
||||
# actual model tokenizer rather than assuming that one character is one token.
|
||||
windows, offset = [], 0
|
||||
while offset < len(text):
|
||||
remaining = text[offset:]
|
||||
if fits(remaining):
|
||||
windows.append(remaining)
|
||||
break
|
||||
low, high = 1, min(len(remaining), maximum * 4)
|
||||
while low < high:
|
||||
middle = (low + high + 1) // 2
|
||||
if fits(remaining[:middle]):
|
||||
low = middle
|
||||
else:
|
||||
high = middle - 1
|
||||
if not fits(remaining[:low]):
|
||||
raise ValueError("单个字符超过向量模型输入上限")
|
||||
windows.append(remaining[:low])
|
||||
offset += low
|
||||
return windows
|
||||
|
||||
|
||||
def embed_texts(texts: list[str]) -> list[list[float]]:
|
||||
if not texts:
|
||||
return []
|
||||
# ``get_encoder`` includes model download and construction only on the
|
||||
# first call in each resident worker. Subsequent calls report zero load
|
||||
# time, which makes accidental worker restarts visible in telemetry.
|
||||
enc = get_encoder()
|
||||
global _encoder_load_ms
|
||||
model_load_ms = _encoder_load_ms
|
||||
# Consume the one-shot load measurement. A resident encoder must report
|
||||
# zero on later requests; a non-zero value then reliably indicates a worker
|
||||
# restart or an explicit model replacement.
|
||||
_encoder_load_ms = 0.0
|
||||
# At most 64 originals and 64 encoded windows in one model batch. Long
|
||||
# paragraphs are pooled in the existing vector space, keeping chunk IDs and
|
||||
# consumption records stable while representing their tails as well.
|
||||
output = []
|
||||
tokenizer_ms = encode_ms = postprocess_ms = 0.0
|
||||
flat_count = 0
|
||||
for start in range(0, len(texts), 64):
|
||||
tokenizer_started = time.perf_counter()
|
||||
groups = [text_windows(str(text), enc) for text in texts[start:start + 64]]
|
||||
flat = [window for group in groups for window in group]
|
||||
tokenizer_ms += (time.perf_counter() - tokenizer_started) * 1000
|
||||
flat_count += len(flat)
|
||||
encode_started = time.perf_counter()
|
||||
vectors = enc.encode(flat, normalize_embeddings=True, show_progress_bar=False, batch_size=64)
|
||||
encode_ms += (time.perf_counter() - encode_started) * 1000
|
||||
postprocess_started = time.perf_counter()
|
||||
cursor = 0
|
||||
for group in groups:
|
||||
selected = vectors[cursor:cursor + len(group)]
|
||||
cursor += len(group)
|
||||
if len(group) == 1:
|
||||
output.append([float(v) for v in selected[0]])
|
||||
continue
|
||||
weights = [max(1, len(enc.tokenizer.encode(window, add_special_tokens=False, truncation=False)))
|
||||
for window in group]
|
||||
combined = [sum(float(vector[i]) * weight for vector, weight in zip(selected, weights))
|
||||
for i in range(len(selected[0]))]
|
||||
norm = math.sqrt(sum(value * value for value in combined)) or 1.0
|
||||
output.append([value / norm for value in combined])
|
||||
postprocess_ms += (time.perf_counter() - postprocess_started) * 1000
|
||||
_set_metric({
|
||||
"model_load_ms": model_load_ms,
|
||||
"tokenizer_ms": tokenizer_ms,
|
||||
"encode_ms": encode_ms,
|
||||
"postprocess_ms": postprocess_ms,
|
||||
"embedding_count": len(texts),
|
||||
# ``flat_count`` is useful when long documents are split into windows;
|
||||
# retain it as a diagnostic without changing the public return value.
|
||||
"window_count": flat_count,
|
||||
"embedding_dimension": len(output[0]) if output else 0,
|
||||
"cache_hit": False,
|
||||
})
|
||||
return output
|
||||
|
||||
|
||||
@lru_cache(maxsize=256)
|
||||
def _cached_query(model: str, policy: str, version: str, dimension: str, text: str) -> tuple[float, ...]:
|
||||
"""Cache only query vectors, keyed by model identity and policy version."""
|
||||
return tuple(embed_texts([text])[0])
|
||||
|
||||
|
||||
def embed_query(text: str) -> list[float]:
|
||||
# lru_cache does not expose whether this invocation hit without comparing
|
||||
# its counters. The comparison is process local and avoids retaining the
|
||||
# query itself in telemetry.
|
||||
before = _cached_query.cache_info()
|
||||
global _embedding_dimension
|
||||
# Keep the cache identity stable across the first call: the actual vector
|
||||
# dimension is discovered only after encoding. Model/version already
|
||||
# identify the encoder contract; ``auto`` prevents the first "unknown"
|
||||
# dimension from creating a second key for the same query.
|
||||
result = list(_cached_query(EMBEDDING_MODEL, EMBEDDING_POLICY,
|
||||
EMBEDDING_MODEL_VERSION, "auto", str(text)))
|
||||
if not _embedding_dimension and result:
|
||||
_embedding_dimension = str(len(result))
|
||||
after = _cached_query.cache_info()
|
||||
cache_hit = after.hits > before.hits
|
||||
metric = last_embedding_metric()
|
||||
metric.update(
|
||||
embedding_count=1,
|
||||
cache_hit=cache_hit,
|
||||
embedding_source="query-vector-lru",
|
||||
model=EMBEDDING_MODEL,
|
||||
model_version=EMBEDDING_MODEL_VERSION,
|
||||
pid=os.getpid(),
|
||||
embedding_dimension=len(result) if result else 0,
|
||||
)
|
||||
if cache_hit:
|
||||
# A cache hit performs no tokenizer, model inference, or serialization;
|
||||
# do not leak timings from the previous (possibly indexing) call.
|
||||
metric.update(model_load_ms=0.0, tokenizer_ms=0.0,
|
||||
encode_ms=0.0, postprocess_ms=0.0, window_count=0)
|
||||
_set_metric(metric)
|
||||
return result
|
||||
@@ -0,0 +1,546 @@
|
||||
"""Local Chinese lexical recall, reciprocal-rank fusion and scoped query caches.
|
||||
|
||||
The keyword channel has its own persisted FTS5 index, so it can find documents
|
||||
outside vector top-k. Ranking is lightweight lexical ranking, not a paid/model
|
||||
reranker. Derived indexes contain only the files in the requested scope.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import OrderedDict
|
||||
from contextlib import contextmanager
|
||||
from copy import deepcopy
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError: # Windows
|
||||
fcntl = None
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
from pathlib import Path
|
||||
import re
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
from ..config import DATA_DIR
|
||||
from . import store
|
||||
|
||||
logger = logging.getLogger("yawen.rag.hybrid")
|
||||
INDEX_DIR = DATA_DIR / "rag_hybrid"
|
||||
VERSION = "zh-bigram-fts5-rrf-v2"
|
||||
CACHE_TTL = 300
|
||||
CACHE_MAX = 192
|
||||
CACHE_MAX_CHARS = 8_000_000
|
||||
_cache: OrderedDict = OrderedDict()
|
||||
_cache_lock = threading.RLock()
|
||||
_build_locks: dict[str, threading.Lock] = {}
|
||||
# Index construction is deliberately kept out of request threads. A single
|
||||
# process can have several libraries become ready at once (for example after a
|
||||
# bulk import), so use a small bounded pool and coalesce duplicate requests by
|
||||
# revision. The persisted SQLite index remains protected by _build_lock.
|
||||
_prewarm_executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="rag-prewarm")
|
||||
_prewarm_state_lock = threading.RLock()
|
||||
_prewarm_pending: set[tuple[int, str]] = set()
|
||||
_lexical_metric_local = threading.local()
|
||||
_fusion_metric_local = threading.local()
|
||||
_fts_pool_lock = threading.RLock()
|
||||
_fts_connection_pool: OrderedDict[str, tuple[tuple, sqlite3.Connection, threading.RLock]] = OrderedDict()
|
||||
_FTS_CONNECTION_MAX = 32
|
||||
_STOP = {"如何", "怎么", "什么", "为什么", "为何", "可以", "需要", "我们", "你们", "他们", "这个", "那个", "哪些", "以及", "一个", "进行", "方法"}
|
||||
_RUNS = re.compile(r"[\u3400-\u9fff]+|[a-z0-9]+(?:[._-][a-z0-9]+)*", re.I)
|
||||
|
||||
|
||||
def tokens(value: str) -> list[str]:
|
||||
"""Overlapping Chinese 2/3-grams retain names/numbers without a dictionary."""
|
||||
result = []
|
||||
for run in _RUNS.findall(str(value or "").lower()):
|
||||
if '\u3400' <= run[0] <= '\u9fff':
|
||||
if len(run) == 1:
|
||||
result.append(run)
|
||||
for size in (2, 3):
|
||||
result.extend(run[i:i + size] for i in range(len(run) - size + 1)
|
||||
if run[i:i + size] not in _STOP)
|
||||
else:
|
||||
result.append(run)
|
||||
return result
|
||||
|
||||
|
||||
def _digest(value) -> str:
|
||||
return hashlib.sha256(json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True).encode()).hexdigest()
|
||||
|
||||
|
||||
def revision(kb_id: int) -> str:
|
||||
"""Store revisions are persisted so invalidation also reaches query workers."""
|
||||
return store.kb_revision(int(kb_id))
|
||||
|
||||
|
||||
def warm_state(org: int, kb_id: int, rev: str | None = None,
|
||||
files: set[str] | None = None) -> dict:
|
||||
"""Return a lightweight, read-only FTS warm state for diagnostics."""
|
||||
try:
|
||||
library_id, current = int(kb_id), str(rev or revision(kb_id))
|
||||
ready = index_available(int(org or 0), library_id, current, files)
|
||||
with _prewarm_state_lock:
|
||||
pending = (library_id, current) in _prewarm_pending
|
||||
return {"state": "ready" if ready else ("warming" if pending else "cold"),
|
||||
"ready": bool(ready), "pending": bool(pending),
|
||||
"revision": current}
|
||||
except Exception as exc:
|
||||
return {"state": "error", "ready": False, "pending": False,
|
||||
"error": type(exc).__name__}
|
||||
|
||||
|
||||
def last_lexical_metric() -> dict:
|
||||
return dict(getattr(_lexical_metric_local, "metric", {}) or {})
|
||||
|
||||
|
||||
def last_fusion_metric() -> dict:
|
||||
return dict(getattr(_fusion_metric_local, "metric", {}) or {})
|
||||
|
||||
|
||||
def _scope_key(org: int, kb_id: int, rev: str, files: set[str] | None) -> str:
|
||||
return _digest([VERSION, int(org), int(kb_id), rev, None if files is None else sorted(files)])
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _build_lock(path: Path):
|
||||
with _cache_lock:
|
||||
lock = _build_locks.setdefault(str(path), threading.Lock())
|
||||
with lock:
|
||||
if fcntl is None:
|
||||
yield
|
||||
return
|
||||
with path.with_suffix('.lock').open('a') as handle:
|
||||
fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
|
||||
|
||||
|
||||
def _valid_index(path: Path, signature: str) -> bool:
|
||||
if not path.exists():
|
||||
return False
|
||||
try:
|
||||
with sqlite3.connect(f"file:{path}?mode=ro", uri=True) as db:
|
||||
row = db.execute("SELECT value FROM settings WHERE key='signature'").fetchone()
|
||||
return bool(row and row[0] == signature)
|
||||
except sqlite3.Error:
|
||||
return False
|
||||
|
||||
|
||||
def _fts_connection(path: Path) -> tuple[sqlite3.Connection, threading.RLock, bool]:
|
||||
"""Return a shared read-only FTS connection for ``path``.
|
||||
|
||||
FTS indexes are immutable between atomic replacements. Reusing a
|
||||
connection avoids repeatedly opening the index on the external volume;
|
||||
the inode/size/mtime signature forces a reconnect whenever a new revision
|
||||
replaces the file, so a worker never serves an old index indefinitely.
|
||||
"""
|
||||
stat = path.stat()
|
||||
signature = (str(path), int(getattr(stat, "st_ino", 0)), int(stat.st_size),
|
||||
int(getattr(stat, "st_mtime_ns", int(stat.st_mtime * 1_000_000_000))))
|
||||
key = str(path)
|
||||
with _fts_pool_lock:
|
||||
holder = _fts_connection_pool.get(key)
|
||||
if holder and holder[0] == signature:
|
||||
_fts_connection_pool.move_to_end(key)
|
||||
return holder[1], holder[2], True
|
||||
# A revision is published with an atomic rename. Serialize rotation
|
||||
# with any query using the previous handle before closing it.
|
||||
gate = holder[2] if holder else threading.RLock()
|
||||
with gate:
|
||||
uri = f"file:{path}?mode=ro"
|
||||
# The bounded recall pool may hand the same KB to different
|
||||
# workers across requests. A per-handle gate serializes access;
|
||||
# disable sqlite's thread identity check so the handle can be
|
||||
# safely reused by those workers.
|
||||
connection = sqlite3.connect(uri, uri=True, check_same_thread=False)
|
||||
# These indexes are read-only and query-heavy. A larger per-
|
||||
# connection page cache and mmap keep repeated Chinese MATCH
|
||||
# queries off the external volume.
|
||||
for pragma in (
|
||||
"PRAGMA query_only=ON",
|
||||
"PRAGMA temp_store=MEMORY",
|
||||
"PRAGMA cache_size=-32768",
|
||||
"PRAGMA mmap_size=268435456",
|
||||
):
|
||||
try:
|
||||
connection.execute(pragma)
|
||||
except sqlite3.Error:
|
||||
# Older SQLite builds may reject mmap_size; the query
|
||||
# remains correct and still benefits from other settings.
|
||||
if pragma != "PRAGMA mmap_size=268435456":
|
||||
connection.close()
|
||||
raise
|
||||
if holder:
|
||||
try:
|
||||
holder[1].close()
|
||||
except sqlite3.Error:
|
||||
pass
|
||||
_fts_connection_pool[key] = (signature, connection, gate)
|
||||
_fts_connection_pool.move_to_end(key)
|
||||
while len(_fts_connection_pool) > _FTS_CONNECTION_MAX:
|
||||
old_key, (_, old_connection, old_gate) = _fts_connection_pool.popitem(last=False)
|
||||
if old_key == key:
|
||||
# Defensive guard; the just-created handle should never
|
||||
# be evicted, but restore it if the limit is misconfigured.
|
||||
_fts_connection_pool[key] = (signature, connection, gate)
|
||||
break
|
||||
with old_gate:
|
||||
try:
|
||||
old_connection.close()
|
||||
except sqlite3.Error:
|
||||
pass
|
||||
return connection, gate, False
|
||||
|
||||
|
||||
def _drop_fts_connection() -> None:
|
||||
# The failed query may have come from any recall worker. Drop all handles
|
||||
# so the next request reopens the affected index with a clean connection.
|
||||
with _fts_pool_lock:
|
||||
for _, connection, connection_gate in list(_fts_connection_pool.values()):
|
||||
with connection_gate:
|
||||
try:
|
||||
connection.close()
|
||||
except sqlite3.Error:
|
||||
pass
|
||||
_fts_connection_pool.clear()
|
||||
|
||||
|
||||
def _index(org: int, kb_id: int, rev: str, files: set[str] | None) -> Path:
|
||||
signature = _scope_key(org, kb_id, rev, files)
|
||||
INDEX_DIR.mkdir(parents=True, exist_ok=True)
|
||||
# One derived file per scope; revisions atomically replace it rather than
|
||||
# retaining a full text copy of every historical library revision.
|
||||
scope_name = _scope_key(org, kb_id, "", files)
|
||||
path = INDEX_DIR / f"kb_{kb_id}_{scope_name}.sqlite3"
|
||||
if _valid_index(path, signature):
|
||||
return path
|
||||
with _build_lock(path):
|
||||
if _valid_index(path, signature):
|
||||
return path
|
||||
stage = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
|
||||
started = time.perf_counter()
|
||||
count = 0
|
||||
try:
|
||||
with sqlite3.connect(str(stage)) as db:
|
||||
db.execute('PRAGMA journal_mode=OFF')
|
||||
db.execute('PRAGMA synchronous=OFF')
|
||||
db.execute('CREATE TABLE settings(key TEXT PRIMARY KEY, value TEXT)')
|
||||
db.execute('CREATE TABLE sources(rowid INTEGER PRIMARY KEY, chunk_id TEXT NOT NULL, source_text TEXT NOT NULL, metadata TEXT NOT NULL)')
|
||||
# Contentless FTS stores postings, not another expanded n-gram copy
|
||||
# of every source document. The original text lives once in sources.
|
||||
db.execute('CREATE VIRTUAL TABLE docs USING fts5(body, heading, content="", tokenize="unicode61")')
|
||||
batch = []
|
||||
for item in store.iter_chunks(kb_id, file_names=files):
|
||||
meta = item.get('meta') or {}
|
||||
# Defense in depth, including future store implementations.
|
||||
if files is not None and meta.get('file') not in files:
|
||||
raise RuntimeError('关键词索引读取到了所选分类以外的文件')
|
||||
source = str(item.get('text') or '')
|
||||
if not source.strip() or not item.get('id'):
|
||||
continue
|
||||
count += 1
|
||||
batch.append((count, ' '.join(tokens(source)), ' '.join(tokens(meta.get('heading', ''))),
|
||||
str(item['id']), source, json.dumps(meta, ensure_ascii=False)))
|
||||
if len(batch) >= 128:
|
||||
db.executemany('INSERT INTO docs(rowid,body,heading) VALUES(?,?,?)', (row[:3] for row in batch))
|
||||
db.executemany('INSERT INTO sources VALUES(?,?,?,?)', ((row[0], *row[3:]) for row in batch))
|
||||
batch.clear()
|
||||
if batch:
|
||||
db.executemany('INSERT INTO docs(rowid,body,heading) VALUES(?,?,?)', (row[:3] for row in batch))
|
||||
db.executemany('INSERT INTO sources VALUES(?,?,?,?)', ((row[0], *row[3:]) for row in batch))
|
||||
if revision(kb_id) != rev:
|
||||
raise RuntimeError('知识库索引正在更新，请稍后重新检索')
|
||||
db.execute('INSERT INTO settings VALUES(?,?)', ('signature', signature))
|
||||
db.execute("INSERT INTO docs(docs) VALUES('optimize')")
|
||||
stage.replace(path)
|
||||
logger.info('中文关键词索引就绪：知识库 %s，%s 段，%.3f 秒', kb_id, count, time.perf_counter() - started)
|
||||
finally:
|
||||
stage.unlink(missing_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
def _index_path(org: int, kb_id: int, rev: str, files: set[str] | None) -> Path:
|
||||
"""Return the persisted index path without opening or building it."""
|
||||
scope_name = _scope_key(org, kb_id, "", files)
|
||||
INDEX_DIR.mkdir(parents=True, exist_ok=True)
|
||||
return INDEX_DIR / f"kb_{int(kb_id)}_{scope_name}.sqlite3"
|
||||
|
||||
|
||||
def index_available(org: int, kb_id: int, rev: str, files: set[str] | None = None) -> bool:
|
||||
"""Return whether a derived FTS index already exists for this scope.
|
||||
|
||||
Building a first index can require scanning tens of thousands of chunks. It
|
||||
must never happen synchronously on a content generation request: callers can
|
||||
use the vector channel immediately and let ``schedule_prewarm`` finish the
|
||||
lexical index in the background.
|
||||
"""
|
||||
try:
|
||||
signature = _scope_key(int(org or 0), int(kb_id), str(rev), files)
|
||||
path = _index_path(int(org or 0), int(kb_id), str(rev), files)
|
||||
return _valid_index(path, signature)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def schedule_prewarm(org: int, kb_id: int, rev: str | None = None) -> bool:
|
||||
"""Build the full-library lexical index in the background.
|
||||
|
||||
Upload/rebuild callers invoke this only after their transaction has
|
||||
committed. It is intentionally best-effort: a failed prewarm must never
|
||||
make an otherwise successful upload fail; the normal first query can still
|
||||
build the index under the same cross-process lock. Requests for the same
|
||||
library revision are coalesced so repeated progress callbacks do not spawn
|
||||
unbounded workers.
|
||||
"""
|
||||
try:
|
||||
library_id = int(kb_id)
|
||||
current_revision = str(rev or revision(library_id))
|
||||
key = (library_id, current_revision)
|
||||
except Exception:
|
||||
return False
|
||||
with _prewarm_state_lock:
|
||||
if key in _prewarm_pending:
|
||||
return False
|
||||
_prewarm_pending.add(key)
|
||||
|
||||
def _worker():
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
# A newer upload may have committed while this task was queued;
|
||||
# skip stale work and let the newer revision be scheduled instead.
|
||||
if revision(library_id) != current_revision:
|
||||
return
|
||||
_index(int(org or 0), library_id, current_revision, None)
|
||||
logger.info("知识库全文索引已后台预热：知识库 %s，%.3f 秒",
|
||||
library_id, time.perf_counter() - started)
|
||||
except Exception:
|
||||
# Best effort only. Keep the exception in server logs for
|
||||
# diagnosis, while preserving the successful import/rebuild API.
|
||||
logger.exception("知识库全文索引后台预热失败：知识库 %s", library_id)
|
||||
finally:
|
||||
with _prewarm_state_lock:
|
||||
_prewarm_pending.discard(key)
|
||||
|
||||
try:
|
||||
_prewarm_executor.submit(_worker)
|
||||
return True
|
||||
except RuntimeError:
|
||||
# Interpreter shutdown or executor teardown; requests can fall back to
|
||||
# lazy construction on the next search.
|
||||
with _prewarm_state_lock:
|
||||
_prewarm_pending.discard(key)
|
||||
return False
|
||||
|
||||
|
||||
def lexical_search(org: int, kb_id: int, rev: str, query: str, limit: int,
|
||||
files: set[str] | None = None) -> list[dict]:
|
||||
started = time.perf_counter()
|
||||
metric = {"kb_id": int(kb_id), "fts_ms": 0.0, "metadata_filter_ms": 0.0,
|
||||
"warm_state": warm_state(org, kb_id, rev, files)}
|
||||
_lexical_metric_local.metric = metric
|
||||
if files is not None and not files:
|
||||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||||
return []
|
||||
terms = list(dict.fromkeys(tokens(query)))[:192]
|
||||
if not terms:
|
||||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||||
return []
|
||||
# A category/file scope can be very large (tens of thousands of chunks).
|
||||
# Building a dedicated FTS database synchronously in the generation
|
||||
# request used to block the RAG worker until its 45s deadline. Full
|
||||
# library indexes are prewarmed after import, so use one as a safe fallback
|
||||
# and filter its metadata after the match. The vector channel still runs
|
||||
# concurrently when a full index is being prepared.
|
||||
post_filter_files = None
|
||||
scope_path = _index_path(org, kb_id, rev, files)
|
||||
scope_signature = _scope_key(org, kb_id, rev, files)
|
||||
if files is not None and not _valid_index(scope_path, scope_signature):
|
||||
full_path = _index_path(org, kb_id, rev, None)
|
||||
full_signature = _scope_key(org, kb_id, rev, None)
|
||||
if not _valid_index(full_path, full_signature):
|
||||
# Do not construct a potentially multi-hundred-MB index on the
|
||||
# request thread. Import/rebuild schedules this work in the
|
||||
# background; returning no lexical hits is preferable to timing
|
||||
# out the whole generation request.
|
||||
schedule_prewarm(org, kb_id, rev)
|
||||
logger.info("知识库 %s 的全文关键词索引尚未就绪，已后台预热；本次仅使用向量召回", kb_id)
|
||||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||||
return []
|
||||
path = full_path
|
||||
post_filter_files = files
|
||||
else:
|
||||
# The full-library index is normally ready from the post-import
|
||||
# prewarm. If a revision changed and the replacement is still being
|
||||
# built, keep this request bounded as well instead of synchronously
|
||||
# scanning a large collection in the RAG worker.
|
||||
if files is None:
|
||||
full_path = _index_path(org, kb_id, rev, None)
|
||||
full_signature = _scope_key(org, kb_id, rev, None)
|
||||
if not _valid_index(full_path, full_signature):
|
||||
schedule_prewarm(org, kb_id, rev)
|
||||
logger.info("知识库 %s 的全文关键词索引尚未就绪，已后台预热；本次仅使用向量召回", kb_id)
|
||||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||||
return []
|
||||
path = full_path
|
||||
else:
|
||||
path = scope_path
|
||||
# Every token is quoted; user input cannot introduce FTS operators.
|
||||
match = ' OR '.join('"' + term.replace('"', '""') + '"' for term in terms)
|
||||
try:
|
||||
fts_started = time.perf_counter()
|
||||
db, db_gate, reused = _fts_connection(path)
|
||||
metric["connection_reused"] = bool(reused)
|
||||
# Pull a bounded oversample when filtering a full-library index so
|
||||
# that relevant files are not lost behind unrelated matches.
|
||||
fetch_limit = min(2000, max(limit, limit * 8)) if post_filter_files is not None else limit
|
||||
with db_gate:
|
||||
rows = db.execute('SELECT sources.chunk_id,sources.source_text,sources.metadata,bm25(docs,1.0,2.2) AS rank FROM docs JOIN sources ON sources.rowid=docs.rowid WHERE docs MATCH ? ORDER BY rank,sources.chunk_id LIMIT ?',
|
||||
(match, fetch_limit)).fetchall()
|
||||
metric["fts_ms"] = round((time.perf_counter() - fts_started) * 1000, 2)
|
||||
except sqlite3.Error as exc:
|
||||
_drop_fts_connection()
|
||||
raise RuntimeError(f'知识库 {kb_id} 关键词检索失败：{exc}') from exc
|
||||
result = []
|
||||
for row in rows:
|
||||
filter_started = time.perf_counter()
|
||||
metadata = json.loads(row[2])
|
||||
if post_filter_files is not None and metadata.get('file') not in post_filter_files:
|
||||
metric["metadata_filter_ms"] += (time.perf_counter() - filter_started) * 1000
|
||||
continue
|
||||
result.append({'id': row[0], 'text': row[1], 'meta': metadata,
|
||||
'lexical_score': max(0.0, -float(row[3]))})
|
||||
if len(result) >= limit:
|
||||
metric["metadata_filter_ms"] += (time.perf_counter() - filter_started) * 1000
|
||||
break
|
||||
metric["metadata_filter_ms"] += (time.perf_counter() - filter_started) * 1000
|
||||
metric["metadata_filter_ms"] = round(metric["metadata_filter_ms"], 2)
|
||||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||||
return result
|
||||
|
||||
|
||||
def cache_get(key):
|
||||
with _cache_lock:
|
||||
cached = _cache.get(key)
|
||||
if cached is None:
|
||||
return None
|
||||
if time.monotonic() - cached[0] > CACHE_TTL:
|
||||
_cache.pop(key, None)
|
||||
return None
|
||||
_cache.move_to_end(key)
|
||||
return deepcopy(cached[1])
|
||||
|
||||
|
||||
def cache_put(key, value):
|
||||
with _cache_lock:
|
||||
size = len(json.dumps(value, ensure_ascii=False))
|
||||
if size > CACHE_MAX_CHARS:
|
||||
return
|
||||
_cache[key] = (time.monotonic(), deepcopy(value), size)
|
||||
_cache.move_to_end(key)
|
||||
while len(_cache) > CACHE_MAX or sum(row[2] for row in _cache.values()) > CACHE_MAX_CHARS:
|
||||
_cache.popitem(last=False)
|
||||
|
||||
|
||||
def clear_cache():
|
||||
with _cache_lock:
|
||||
_cache.clear()
|
||||
|
||||
|
||||
def prewarm_index(org: int, kb_id: int, files: set[str] | None = None) -> None:
|
||||
"""Build the FTS scope outside the user's request path.
|
||||
|
||||
Import/rebuild jobs call this after the vector index is committed. The
|
||||
existing file lock makes concurrent warmups idempotent and safe across
|
||||
API and worker processes.
|
||||
"""
|
||||
_index(int(org or 1), int(kb_id), revision(int(kb_id)), files)
|
||||
|
||||
|
||||
def candidate_key(org: int, user: int, scopes: list[tuple], query: str, limit: int,
|
||||
model: str, *, mode: str = "hybrid", top_k: int | None = None,
|
||||
permissions=None, index_version: str | None = None) -> str:
|
||||
"""Build an isolated retrieval candidate-pool cache key.
|
||||
|
||||
Query vectors are cached in :mod:`embed` and deliberately do not use this
|
||||
key. Candidate pools depend on every scope and authorization dimension;
|
||||
callers should pass a stable permission signature when one is available.
|
||||
Defaults preserve compatibility with older six-argument callers.
|
||||
"""
|
||||
# ``scopes`` already carries KB id, index revision and category/file scope.
|
||||
# Sorting permission values makes equivalent grant sets map to one key while
|
||||
# still preventing cross-user or cross-permission pool reuse.
|
||||
permission_scope = permissions
|
||||
if isinstance(permission_scope, (set, frozenset, list, tuple)):
|
||||
permission_scope = sorted(str(value) for value in permission_scope)
|
||||
return _digest([VERSION, int(org or 0), int(user or 0), scopes,
|
||||
str(query).strip(), int(limit), int(top_k or limit),
|
||||
str(model), str(mode or "hybrid"), permission_scope,
|
||||
str(index_version or "")])
|
||||
|
||||
|
||||
def fuse(vector: list[dict], lexical: list[dict], query: str) -> list[dict]:
|
||||
"""RRF balances independent recalls; token coverage refines bounded candidates."""
|
||||
started = time.perf_counter()
|
||||
metric = {"merge_ms": 0.0, "dedup_ms": 0.0, "rrf_ms": 0.0}
|
||||
_fusion_metric_local.metric = metric
|
||||
merge_started = time.perf_counter()
|
||||
merged = {}
|
||||
for channel, rows in (('vector', vector), ('lexical', lexical)):
|
||||
for rank, row in enumerate(rows, 1):
|
||||
key = (int(row['kb_id']), str(row['id']))
|
||||
current = merged.setdefault(key, {**deepcopy(row), 'retrieval_channels': [], '_rrf': 0.0})
|
||||
current['retrieval_channels'].append(channel)
|
||||
current['_rrf'] += 1.0 / (60 + rank)
|
||||
if channel == 'lexical':
|
||||
current['lexical_score'] = row.get('lexical_score', 0)
|
||||
metric["merge_ms"] = round((time.perf_counter() - merge_started) * 1000, 2)
|
||||
metric["dedup_ms"] = metric["merge_ms"]
|
||||
rrf_started = time.perf_counter()
|
||||
qterms = set(tokens(query))
|
||||
# Candidate-local inverse document frequency makes rare, exact names useful.
|
||||
df = {}
|
||||
term_sets = {}
|
||||
for key, row in merged.items():
|
||||
terms = set(tokens(row.get('text', '') + ' ' + str((row.get('meta') or {}).get('heading', ''))))
|
||||
term_sets[key] = terms
|
||||
for term in qterms & terms:
|
||||
df[term] = df.get(term, 0) + 1
|
||||
weights = {term: math.log(1 + (len(merged) + 1) / (df.get(term, 0) + 1)) for term in qterms}
|
||||
total_weight = sum(weights.values()) or 1
|
||||
for key, row in merged.items():
|
||||
coverage = sum(weights[term] for term in qterms & term_sets[key]) / total_weight
|
||||
row['score'] = row.pop('_rrf') + 0.045 * coverage
|
||||
row['keyword_coverage'] = round(coverage, 4)
|
||||
result = sorted(merged.values(), key=lambda row: (-row['score'], int(row['kb_id']), str(row['id'])))
|
||||
metric["rrf_ms"] = round((time.perf_counter() - rrf_started) * 1000, 2)
|
||||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||||
return result
|
||||
|
||||
|
||||
def finalize(candidates: list[dict], k: int, excluded: dict[int, set[str]],
|
||||
allowed: dict[int, set[str] | None]) -> list[dict]:
|
||||
"""Always run after cache lookup: consumption and access scopes are live."""
|
||||
output, seen = [], set()
|
||||
for row in candidates:
|
||||
kid, meta = int(row['kb_id']), row.get('meta') or {}
|
||||
if kid not in allowed or row['id'] in excluded.get(kid, set()):
|
||||
continue
|
||||
files = allowed[kid]
|
||||
if files is not None and meta.get('file') not in files:
|
||||
continue
|
||||
# Same heading/body in one KB is a repeat; numbers and distinct headings survive.
|
||||
fingerprint = (kid, re.sub(r'\s+', ' ', str(meta.get('heading') or '')).strip(),
|
||||
re.sub(r'\s+', ' ', row.get('text', '')).strip())
|
||||
if fingerprint in seen:
|
||||
continue
|
||||
seen.add(fingerprint)
|
||||
output.append(row)
|
||||
if len(output) >= k:
|
||||
break
|
||||
return output
|
||||
@@ -0,0 +1,100 @@
|
||||
"""音视频转写：faster-whisper 本地模型 + ffmpeg 抽音轨。
|
||||
输出与 .srt 同格式的「去时间戳、按语义间隙合并」文本，直接进入既有 RAG 分块管线。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from .parse import clean_text
|
||||
|
||||
MEDIA_SUFFIXES = {".mp4", ".mov", ".m4a", ".mp3", ".wav", ".flv", ".mkv", ".avi", ".webm", ".aac", ".ogg", ".flac"}
|
||||
|
||||
# 模型目录：优先环境变量，其次 modelscope 缓存里的 gpustack/faster-whisper-small
|
||||
_MODEL_DIR_CANDIDATES = [
|
||||
os.environ.get("WHISPER_MODEL_DIR", ""),
|
||||
str(Path.home() / ".cache/modelscope/models/gpustack--faster-whisper-small/snapshots/master"),
|
||||
]
|
||||
_model = None
|
||||
_model_lock = threading.Lock()
|
||||
|
||||
|
||||
def _get_model():
|
||||
global _model
|
||||
if _model is not None:
|
||||
return _model
|
||||
with _model_lock:
|
||||
if _model is not None:
|
||||
return _model
|
||||
from faster_whisper import WhisperModel
|
||||
for cand in _MODEL_DIR_CANDIDATES:
|
||||
if cand and Path(cand).exists():
|
||||
_model = WhisperModel(cand, device="cpu", compute_type="int8", local_files_only=True)
|
||||
return _model
|
||||
_model = WhisperModel("small", device="cpu", compute_type="int8")
|
||||
return _model
|
||||
|
||||
|
||||
def is_media(path: Path) -> bool:
|
||||
return path.suffix.lower() in MEDIA_SUFFIXES
|
||||
|
||||
|
||||
def _probe_duration(path: Path) -> float:
|
||||
try:
|
||||
r = subprocess.run(
|
||||
["ffprobe", "-v", "quiet", "-show_entries", "format=duration",
|
||||
"-of", "csv=p=0", str(path)],
|
||||
capture_output=True, text=True, timeout=30)
|
||||
return float(r.stdout.strip() or 0)
|
||||
except Exception:
|
||||
return 0.0
|
||||
|
||||
|
||||
def transcribe_media(path: Path, progress_cb=None) -> list[dict]:
|
||||
"""音视频 → 语义段落列表 [{text, meta:{heading}}]。
|
||||
超长音先压缩码率降体积；VAD 过滤静音；按标点+间隙合并成语义段。"""
|
||||
src = Path(path)
|
||||
work = src
|
||||
# 大文件先转 16k mono wav（whisper 输入要求，同时把视频降为纯音轨）
|
||||
if src.suffix.lower() not in (".wav",) or _probe_duration(src) > 0:
|
||||
work = src.with_suffix(".wav.wav")
|
||||
cmd = ["ffmpeg", "-y", "-i", str(src), "-vn", "-ar", "16000", "-ac", "1",
|
||||
"-b:a", "64k", str(work)]
|
||||
r = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
|
||||
if r.returncode != 0 or not work.exists() or work.stat().st_size == 0:
|
||||
raise ValueError(f"ffmpeg 抽音轨失败: {r.stderr[-200:]}")
|
||||
try:
|
||||
model = _get_model()
|
||||
segments, _info = model.transcribe(str(work), language="zh", vad_filter=True,
|
||||
vad_parameters={"min_silence_duration_ms": 700})
|
||||
lines: list[str] = []
|
||||
for seg in segments:
|
||||
t = (seg.text or "").strip()
|
||||
if t:
|
||||
lines.append(t)
|
||||
if progress_cb and len(lines) % 20 == 0:
|
||||
progress_cb(len(lines))
|
||||
if not lines:
|
||||
return []
|
||||
# 合并成语义段：约 15 句或 1200 字一段（与既有 2200 字分块目标匹配）
|
||||
merged, buf, size = [], [], 0
|
||||
for t in lines:
|
||||
buf.append(t)
|
||||
size += len(t)
|
||||
if size >= 1200 or t.endswith(("。", "！", "？", "!", "?")) and size >= 600:
|
||||
merged.append(" ".join(buf))
|
||||
buf, size = [], 0
|
||||
if buf:
|
||||
merged.append(" ".join(buf))
|
||||
# 转写结果也经过同一套结构清洗：去掉字幕式序号/时间戳、不可见字符和句首独立口头填充词，
|
||||
# 保留正文中的日期、金额、型号等数字。
|
||||
return [{"text": clean_text(m), "meta": {"heading": src.stem}}
|
||||
for m in merged if clean_text(m)]
|
||||
finally:
|
||||
if work != src:
|
||||
try:
|
||||
work.unlink(missing_ok=True)
|
||||
except OSError:
|
||||
pass
|
||||
@@ -0,0 +1,182 @@
|
||||
"""RAG 管道：文件解析"""
|
||||
import json
|
||||
import re
|
||||
import unicodedata
|
||||
from pathlib import Path
|
||||
|
||||
# 口语稿和字幕里常见的填充词。只在句子边界作为独立词出现时删除，
|
||||
# 避免把“然后我们……”这类有实际语义的连接词误删。
|
||||
_FILLER_WORDS = r"嗯|呃|额|啊|哦|唉|那个|这个嘛|就是嘛"
|
||||
_FILLER_TOKEN_RE = re.compile(
|
||||
rf"(?:(?<=^)|(?<=[\s，。！？；：,.!?;:、]))(?:{_FILLER_WORDS})+(?=(?:[\s，。！？；：,.!?;:、]|$))"
|
||||
)
|
||||
_TIMESTAMP_RE = re.compile(
|
||||
r"^(?:\d{1,4}|\[?\d{1,2}:\d{2}(?::\d{2})?(?:[,.]\d{1,3})?\]?"
|
||||
r"|\d{1,2}:\d{2}(?::\d{2})?(?:[,.]\d{1,3})?\s*[-–—>]\s*.+)$"
|
||||
)
|
||||
|
||||
|
||||
def clean_text(value: str, *, remove_fillers: bool = True) -> str:
|
||||
"""清理导入素材中的格式噪声，同时保留有业务意义的数字。
|
||||
|
||||
只删除整行序号/时间戳、不可见控制字符、重复空白和句末重复标点；
|
||||
数字出现在正文中（如价格、日期、型号）时会原样保留。知识库默认去除
|
||||
句首独立口头填充词；数字人蒸馏可传 ``remove_fillers=False`` 保留口吻。
|
||||
"""
|
||||
if not isinstance(value, str):
|
||||
value = str(value or "")
|
||||
text = unicodedata.normalize("NFKC", value)
|
||||
text = text.replace("\uFEFF", "").replace("\u200B", "").replace("\u200C", "").replace("\u200D", "")
|
||||
# 保留换行和制表，移除其它 C0/C1 控制字符。
|
||||
text = "".join(ch for ch in text if ch in "\n\r\t" or not unicodedata.category(ch).startswith("C"))
|
||||
lines: list[str] = []
|
||||
for raw_line in text.replace("\r\n", "\n").replace("\r", "\n").split("\n"):
|
||||
line = re.sub(r"[ \t]+", " ", raw_line).strip()
|
||||
if not line:
|
||||
continue
|
||||
# 字幕编号、单独页码、时间戳行是结构噪声；正文中的数字不受影响。
|
||||
if _TIMESTAMP_RE.fullmatch(line):
|
||||
continue
|
||||
if remove_fillers:
|
||||
line = _FILLER_TOKEN_RE.sub("", line)
|
||||
line = re.sub(r"^[、,，;；:：\s]+", "", line)
|
||||
# 仅压缩明显的重复标点，不改写省略号和正文标点风格。
|
||||
line = re.sub(r"([。！？!?；;])\1+", r"\1", line)
|
||||
if line:
|
||||
lines.append(line)
|
||||
return "\n".join(lines).strip()
|
||||
|
||||
def parse_file(path: Path) -> list[dict]:
|
||||
"""返回 [{text, meta:{...}}] 按文件内章节/段落拆分的最初单元"""
|
||||
ext = path.suffix.lower()
|
||||
if ext == ".docx":
|
||||
return _parse_docx(path)
|
||||
if ext == ".srt":
|
||||
return _parse_srt(path)
|
||||
if ext == ".pdf":
|
||||
return _parse_pdf(path)
|
||||
if ext == ".json" or ext == ".jsonl":
|
||||
return _parse_json(path)
|
||||
if ext == ".md":
|
||||
return _parse_text(path, is_md=True)
|
||||
if ext in (".txt", ""):
|
||||
return _parse_text(path)
|
||||
return []
|
||||
|
||||
def _parse_pdf(path: Path) -> list[dict]:
|
||||
"""PDF 文本层解析；扫描版无文本时返回空，避免把页码等结构噪声写入索引。"""
|
||||
try:
|
||||
from pypdf import PdfReader
|
||||
reader = PdfReader(str(path))
|
||||
except Exception:
|
||||
return []
|
||||
out: list[dict] = []
|
||||
for page_no, page in enumerate(reader.pages, 1):
|
||||
text = clean_text(page.extract_text() or "")
|
||||
if text:
|
||||
out.append({"text": text, "meta": {"page": page_no}})
|
||||
return out
|
||||
|
||||
def _parse_docx(path: Path) -> list[dict]:
|
||||
try:
|
||||
from docx import Document
|
||||
except ImportError:
|
||||
return []
|
||||
doc = Document(str(path))
|
||||
out, buf, cur_h = [], [], ""
|
||||
for p in doc.paragraphs:
|
||||
t = clean_text(p.text)
|
||||
if not t:
|
||||
continue
|
||||
if p.style and p.style.name and "Heading" in p.style.name:
|
||||
if buf:
|
||||
out.append({"text": clean_text("\n".join(buf)), "meta": {"heading": cur_h}})
|
||||
buf = []
|
||||
cur_h = t
|
||||
else:
|
||||
buf.append(t)
|
||||
if buf:
|
||||
out.append({"text": clean_text("\n".join(buf)), "meta": {"heading": cur_h}})
|
||||
return out
|
||||
|
||||
def _parse_srt(path: Path) -> list[dict]:
|
||||
"""SRT：去时间戳，按时间戳间隙合并成语义段（间隙>2秒 或 字幕号跳变）"""
|
||||
raw = path.read_text(encoding="utf-8", errors="ignore")
|
||||
blocks = re.split(r"\n\s*\n", raw)
|
||||
cur, cur_end, out = [], None, []
|
||||
for b in blocks:
|
||||
lines = [l.strip() for l in b.splitlines() if l.strip()]
|
||||
if not lines:
|
||||
continue
|
||||
# 行1=序号，行2=时间 00:00:01,000 --> 00:00:03,500
|
||||
if len(lines) >= 2 and "-->" in lines[1]:
|
||||
m = re.match(r"(\d+):(\d+):(\d+),(\d+)", lines[1])
|
||||
start = int(m.group(1)) * 3600 + int(m.group(2)) * 60 + int(m.group(3)) if m else 0
|
||||
gap = (start - cur_end) if cur_end is not None else 0
|
||||
text = clean_text(" ".join(lines[2:]))
|
||||
if cur and gap > 2.0:
|
||||
out.append({"text": clean_text(" ".join(cur)), "meta": {}})
|
||||
cur = []
|
||||
cur.append(text)
|
||||
cur_end = start
|
||||
if cur:
|
||||
out.append({"text": clean_text(" ".join(cur)), "meta": {}})
|
||||
return out
|
||||
|
||||
def _parse_json(path: Path) -> list[dict]:
|
||||
raw = path.read_text(encoding="utf-8", errors="ignore")
|
||||
out = []
|
||||
# JSONL：逐行 JSON（观品 chunks.jsonl 格式）
|
||||
for line in raw.splitlines():
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
item = json.loads(line)
|
||||
except Exception:
|
||||
continue
|
||||
if isinstance(item, dict):
|
||||
text = item.get("content") or item.get("text") or item.get("page_content")
|
||||
if text:
|
||||
text = clean_text(str(text))
|
||||
if " " in text:
|
||||
text = text.replace(" ", "") # 观品格式字间空格
|
||||
out.append({"text": text, "meta": {"heading": item.get("title", "")}})
|
||||
if out:
|
||||
return out
|
||||
data = json.loads(raw)
|
||||
if isinstance(data, list):
|
||||
for i, item in enumerate(data):
|
||||
text = item.get("content") or item.get("text") or (item.get("page_content") if isinstance(item, dict) else str(item))
|
||||
if text:
|
||||
out.append({"text": clean_text(str(text)), "meta": {"index": i}})
|
||||
elif isinstance(data, dict):
|
||||
# 支持 {title: text} 或 {chunks: [...]}
|
||||
for k, v in data.items():
|
||||
if isinstance(v, str) and v.strip():
|
||||
out.append({"text": clean_text(v), "meta": {"heading": clean_text(k, remove_fillers=False)}})
|
||||
elif isinstance(v, list):
|
||||
for item in v:
|
||||
t = item.get("content") or item.get("text") if isinstance(item, dict) else str(item)
|
||||
if t:
|
||||
out.append({"text": clean_text(str(t)), "meta": {"heading": clean_text(k, remove_fillers=False)}})
|
||||
return out
|
||||
|
||||
def _parse_text(path: Path, is_md: bool = False) -> list[dict]:
|
||||
out, buf, cur_h = [], [], ""
|
||||
# 按行读取，避免公网分片上传后的超大文本一次性进入内存。
|
||||
with path.open("r", encoding="utf-8", errors="ignore") as source:
|
||||
for line in source:
|
||||
s = clean_text(line)
|
||||
if not s:
|
||||
continue
|
||||
if is_md and (s.startswith("# ") or s.startswith("## ")):
|
||||
if buf:
|
||||
out.append({"text": clean_text("\n".join(buf)), "meta": {"heading": cur_h}})
|
||||
buf = []
|
||||
cur_h = s.lstrip("#").strip()
|
||||
else:
|
||||
buf.append(s)
|
||||
if buf:
|
||||
out.append({"text": clean_text("\n".join(buf)), "meta": {"heading": cur_h}})
|
||||
return out
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,309 @@
|
||||
"""Qdrant REST adapter preserving the application's original chunk IDs and metadata.
|
||||
|
||||
The small collection interface matches the store's existing snapshot/rollback code.
|
||||
REST calls have deadlines; backend errors never fall back to a stale local index.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import hashlib
|
||||
import math
|
||||
import threading
|
||||
import uuid
|
||||
import httpx
|
||||
|
||||
|
||||
class QdrantError(RuntimeError):
|
||||
def __init__(self, message, *, status_code=None):
|
||||
super().__init__(message)
|
||||
self.status_code = status_code
|
||||
|
||||
|
||||
class QdrantClient:
|
||||
def __init__(self, url, api_key, prefix='yawan_', *, http_client=None):
|
||||
if not prefix or not all(c.isalnum() or c in '_-' for c in prefix):
|
||||
raise ValueError('Qdrant 知识库命名空间无效')
|
||||
self.prefix = prefix
|
||||
self._creation_lock = threading.Lock()
|
||||
self.http = http_client or httpx.Client(base_url=url.rstrip('/'),
|
||||
headers={'api-key': api_key}, timeout=httpx.Timeout(25, connect=5),
|
||||
limits=httpx.Limits(max_connections=8, max_keepalive_connections=4), trust_env=False)
|
||||
|
||||
def request(self, method, path, body=None, *, missing_ok=False):
|
||||
try:
|
||||
response = self.http.request(method, path, json=body)
|
||||
except httpx.TimeoutException as exc:
|
||||
raise QdrantError('Qdrant 知识索引请求超时，请检查向量服务') from exc
|
||||
except httpx.HTTPError as exc:
|
||||
raise QdrantError('Qdrant 知识索引连接失败，请检查向量服务') from exc
|
||||
if missing_ok and response.status_code == 404:
|
||||
return None
|
||||
if not response.is_success:
|
||||
# No body: it can contain stored user text or reflected credentials.
|
||||
raise QdrantError(f'Qdrant 知识索引操作失败（HTTP {response.status_code}）', status_code=response.status_code)
|
||||
try:
|
||||
payload = response.json()
|
||||
return payload.get('result', payload)
|
||||
except (ValueError, AttributeError) as exc:
|
||||
raise QdrantError('Qdrant 知识索引返回了无效数据') from exc
|
||||
|
||||
def get_collection(self, name):
|
||||
collection = QdrantCollection(self, name)
|
||||
result = self.request('GET', collection.path, missing_ok=True)
|
||||
if result is None:
|
||||
raise QdrantError(f'知识索引集合不存在：{name}')
|
||||
collection.configure(result)
|
||||
return collection
|
||||
|
||||
def get_or_create_collection(self, name):
|
||||
collection = QdrantCollection(self, name)
|
||||
with self._creation_lock:
|
||||
description = self.request('GET', collection.path, missing_ok=True)
|
||||
if description is None:
|
||||
try:
|
||||
self.request('PUT', collection.path, {
|
||||
'vectors': {'size': 512, 'distance': 'Cosine'},
|
||||
'on_disk_payload': True,
|
||||
'optimizers_config': {'default_segment_number': 2},
|
||||
'hnsw_config': {'m': 16, 'ef_construct': 128}})
|
||||
except QdrantError as exc:
|
||||
# Another process can win after our GET; only accept a
|
||||
# confirmed, compatible collection, never hide server errors.
|
||||
if exc.status_code not in {400, 409}:
|
||||
raise
|
||||
description = self.request('GET', collection.path, missing_ok=True)
|
||||
if description is None:
|
||||
raise exc
|
||||
else:
|
||||
description = {'config': {'params': {'vectors': {'size': 512, 'distance': 'Cosine'}}}}
|
||||
collection.configure(description)
|
||||
self.request('PUT', collection.path + '/index?wait=true',
|
||||
{'field_name': 'meta.file', 'field_schema': 'keyword'})
|
||||
else:
|
||||
collection.configure(description)
|
||||
return collection
|
||||
|
||||
def delete_collection(self, name):
|
||||
return self.request('DELETE', QdrantCollection(self, name).path, missing_ok=True)
|
||||
|
||||
def close(self):
|
||||
self.http.close()
|
||||
|
||||
|
||||
class QdrantCollection:
|
||||
def __init__(self, client, name):
|
||||
if not name or not all(c.isalnum() or c in '_-' for c in name):
|
||||
raise ValueError('知识索引集合名称无效')
|
||||
self.client, self.name = client, name
|
||||
self.path = '/collections/' + client.prefix + name
|
||||
self.distance = 'Cosine'
|
||||
|
||||
def configure(self, description):
|
||||
try:
|
||||
vectors = description['config']['params']['vectors']
|
||||
except (KeyError, TypeError):
|
||||
vectors = None
|
||||
if not isinstance(vectors, dict) or vectors.get('size') != 512 or vectors.get('distance') not in {'Cosine', 'Euclid'}:
|
||||
raise QdrantError('知识索引集合必须使用 512 维 Cosine 或 Euclid 向量')
|
||||
self.distance = vectors['distance']
|
||||
|
||||
def point_id(self, original_id):
|
||||
return str(uuid.uuid5(uuid.NAMESPACE_URL, self.client.prefix + self.name + ':' + str(original_id)))
|
||||
|
||||
@staticmethod
|
||||
def filter(where):
|
||||
if not where:
|
||||
return None
|
||||
if not isinstance(where, dict) or set(where) != {'file'}:
|
||||
raise ValueError('不支持的知识索引过滤条件')
|
||||
value = where['file']
|
||||
if isinstance(value, dict):
|
||||
if set(value) != {'$in'} or not isinstance(value['$in'], (list, tuple, set)) or any(not isinstance(v, str) for v in value['$in']):
|
||||
raise ValueError('知识索引文件过滤条件无效')
|
||||
match = {'any': list(value['$in'])}
|
||||
elif isinstance(value, str):
|
||||
match = {'value': value}
|
||||
else:
|
||||
raise ValueError('知识索引文件过滤条件无效')
|
||||
return {'must': [{'key': 'meta.file', 'match': match}]}
|
||||
|
||||
def count(self, where=None):
|
||||
body = {'exact': True}
|
||||
filtering = self.filter(where)
|
||||
if filtering is not None:
|
||||
body['filter'] = filtering
|
||||
result = self.client.request('POST', self.path + '/points/count', body)
|
||||
return int(result['count'])
|
||||
|
||||
def _pages(self, where, include, size):
|
||||
if isinstance(size, bool) or not isinstance(size, int) or not 1 <= size <= 4000:
|
||||
raise ValueError('索引分页大小无效')
|
||||
offset = None
|
||||
seen_offsets = set()
|
||||
while True:
|
||||
body = {'limit': size, 'with_vector': 'embeddings' in include,
|
||||
'with_payload': {'include': ['chunk_id', 'meta'] + (['text'] if 'documents' in include else [])}}
|
||||
filtering = self.filter(where)
|
||||
if filtering is not None:
|
||||
body['filter'] = filtering
|
||||
if offset is not None:
|
||||
body['offset'] = offset
|
||||
result = self.client.request('POST', self.path + '/points/scroll', body)
|
||||
yield result.get('points', [])
|
||||
new_offset = result.get('next_page_offset')
|
||||
if new_offset is None:
|
||||
break
|
||||
if new_offset in seen_offsets:
|
||||
raise QdrantError('Qdrant 分页游标未前进')
|
||||
seen_offsets.add(new_offset)
|
||||
offset = new_offset
|
||||
|
||||
@staticmethod
|
||||
def _format(points, include):
|
||||
result = {'ids': [], 'documents': [], 'metadatas': [], 'embeddings': []}
|
||||
for point in points:
|
||||
payload = point.get('payload') or {}
|
||||
if not isinstance(payload.get('chunk_id'), str) or not payload['chunk_id']:
|
||||
raise QdrantError('知识索引原始片段编号缺失')
|
||||
result['ids'].append(payload['chunk_id'])
|
||||
if 'documents' in include:
|
||||
if not isinstance(payload.get('text'), str):
|
||||
raise QdrantError('知识索引原文缺失')
|
||||
result['documents'].append(payload['text'])
|
||||
if 'metadatas' in include:
|
||||
if not isinstance(payload.get('meta'), dict):
|
||||
raise QdrantError('知识索引来源元数据缺失')
|
||||
result['metadatas'].append(payload['meta'])
|
||||
if 'embeddings' in include:
|
||||
vector = point.get('vector')
|
||||
if not isinstance(vector, list) or len(vector) != 512 or any(not isinstance(v, (int, float)) or not math.isfinite(v) for v in vector):
|
||||
raise QdrantError('知识索引向量维度无效')
|
||||
result['embeddings'].append(vector)
|
||||
return result
|
||||
|
||||
def iter_records(self, where=None, include=None, page_size=256):
|
||||
include = ['documents', 'metadatas'] if include is None else include
|
||||
seen = set()
|
||||
for page in self._pages(where, include, page_size):
|
||||
self._check_scope(page, where)
|
||||
result = self._format(page, include)
|
||||
for identifier in result['ids']:
|
||||
if identifier in seen:
|
||||
raise QdrantError('知识索引分页返回重复片段，请重新读取')
|
||||
seen.add(identifier)
|
||||
yield result
|
||||
|
||||
@staticmethod
|
||||
def _check_scope(points, where):
|
||||
if not where:
|
||||
return
|
||||
allowed = where['file']
|
||||
for point in points:
|
||||
name = (point.get('payload', {}).get('meta') or {}).get('file')
|
||||
if not (name in allowed['$in'] if isinstance(allowed, dict) else name == allowed):
|
||||
raise QdrantError('知识索引返回了文件范围外的片段')
|
||||
|
||||
def get(self, *, ids=None, where=None, include=None, limit=None, offset=0):
|
||||
include = ['documents', 'metadatas'] if include is None else include
|
||||
if isinstance(offset, bool) or not isinstance(offset, int) or offset < 0 or (limit is not None and (isinstance(limit, bool) or not isinstance(limit, int) or limit < 0)):
|
||||
raise ValueError('索引分页范围无效')
|
||||
if limit == 0 or ids == []:
|
||||
return self._format([], include)
|
||||
if ids is not None:
|
||||
self.filter(where)
|
||||
found = {}
|
||||
for start in range(0, len(ids), 128):
|
||||
batch = ids[start:start + 128]
|
||||
points = self.client.request('POST', self.path + '/points',
|
||||
{'ids': [self.point_id(identifier) for identifier in batch],
|
||||
'with_payload': True, 'with_vector': 'embeddings' in include})
|
||||
for point in points:
|
||||
original_id = (point.get('payload') or {}).get('chunk_id')
|
||||
if original_id not in batch:
|
||||
raise QdrantError('知识索引返回了请求范围外的片段')
|
||||
if where:
|
||||
name = (point.get('payload', {}).get('meta') or {}).get('file')
|
||||
allowed = where['file']
|
||||
if not (name in allowed['$in'] if isinstance(allowed, dict) else name == allowed):
|
||||
continue
|
||||
found[original_id] = point
|
||||
selected = [found[identifier] for identifier in ids if identifier in found]
|
||||
return self._format(selected[offset:None if limit is None else offset + limit], include)
|
||||
retained = []
|
||||
visited = 0
|
||||
for page in self._pages(where, include, min(256, max(1, limit or 256))):
|
||||
self._check_scope(page, where)
|
||||
for point in page:
|
||||
if visited >= offset:
|
||||
retained.append(point)
|
||||
visited += 1
|
||||
if limit is not None and len(retained) >= limit:
|
||||
return self._format(retained, include)
|
||||
return self._format(retained, include)
|
||||
|
||||
def add(self, *, ids, documents, embeddings, metadatas):
|
||||
if not (len(ids) == len(documents) == len(embeddings) == len(metadatas)):
|
||||
raise ValueError('索引编号、正文、向量和元数据数量不一致')
|
||||
# Validate the entire input before the first batch is written.
|
||||
if any(not isinstance(identifier, str) or not identifier for identifier in ids) or len(set(ids)) != len(ids):
|
||||
raise ValueError('索引片段编号为空或重复')
|
||||
for text, vector, meta in zip(documents, embeddings, metadatas):
|
||||
if not isinstance(text, str) or not text.strip() or not isinstance(meta, dict):
|
||||
raise ValueError('索引正文或元数据无效')
|
||||
if not isinstance(vector, (list, tuple)) or len(vector) != 512 or any(not isinstance(v, (int, float)) or not math.isfinite(v) for v in vector):
|
||||
raise ValueError('索引向量必须是 512 维有限数值')
|
||||
for start in range(0, len(ids), 128):
|
||||
points = [{'id': self.point_id(ids[i]), 'vector': embeddings[i],
|
||||
'payload': {'chunk_id': ids[i], 'text': documents[i], 'meta': metadatas[i]}}
|
||||
for i in range(start, min(len(ids), start + 128))]
|
||||
self.client.request('PUT', self.path + '/points?wait=true', {'points': points})
|
||||
|
||||
def delete(self, *, where=None, ids=None):
|
||||
if ids is not None and where is not None:
|
||||
raise ValueError('删除知识索引不可同时指定两种范围')
|
||||
if ids is not None:
|
||||
if not ids:
|
||||
return None
|
||||
body = {'points': [self.point_id(identifier) for identifier in ids]}
|
||||
elif where:
|
||||
body = {'filter': self.filter(where)}
|
||||
else:
|
||||
raise ValueError('删除知识索引必须指定文件或片段编号')
|
||||
return self.client.request('POST', self.path + '/points/delete?wait=true', body)
|
||||
|
||||
def query(self, *, query_embeddings, n_results=6, where=None):
|
||||
if isinstance(n_results, bool) or not isinstance(n_results, int) or n_results < 1:
|
||||
raise ValueError('知识索引检索数量无效')
|
||||
result = {'ids': [], 'documents': [], 'metadatas': [], 'distances': []}
|
||||
for vector in query_embeddings:
|
||||
body = {'query': vector, 'limit': n_results, 'with_payload': True,
|
||||
'with_vector': False, 'params': {'hnsw_ef': max(64, min(256, n_results * 2))}}
|
||||
filtering = self.filter(where)
|
||||
if filtering is not None:
|
||||
body['filter'] = filtering
|
||||
points = self.client.request('POST', self.path + '/points/query', body).get('points', [])
|
||||
self._check_scope(points, where)
|
||||
formatted = self._format(points, ['documents', 'metadatas'])
|
||||
for key in ('ids', 'documents', 'metadatas'):
|
||||
result[key].append(formatted[key])
|
||||
# Chroma's existing default is squared L2. Our embeddings are normalized.
|
||||
result['distances'].append([max(0., 2. * (1. - float(point['score']))) if self.distance == 'Cosine'
|
||||
else float(point['score']) ** 2 for point in points])
|
||||
return result
|
||||
|
||||
|
||||
_client = None
|
||||
_client_signature = None
|
||||
_lock = threading.Lock()
|
||||
|
||||
|
||||
def get_client():
|
||||
from .backend import qdrant_settings
|
||||
global _client, _client_signature
|
||||
url, key, prefix = qdrant_settings()
|
||||
signature = (url, hashlib.sha256(key.encode()).hexdigest(), prefix)
|
||||
with _lock:
|
||||
if _client is None or _client_signature != signature:
|
||||
if _client is not None:
|
||||
_client.close()
|
||||
_client = QdrantClient(url, key, prefix)
|
||||
_client_signature = signature
|
||||
return _client
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Reversible exclusions for explicitly confirmed orphan test indexes.
|
||||
|
||||
The manifest is local operational configuration, not an index mutation. Existing
|
||||
registered files always take precedence and automatically release an exclusion.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from sqlmodel import select
|
||||
|
||||
from ..config import DATA_DIR
|
||||
from ..models import KnowledgeFile
|
||||
|
||||
QUARANTINE_PATH = DATA_DIR / 'rag-quarantine.json'
|
||||
|
||||
|
||||
class QuarantineManifestError(ValueError):
|
||||
"""Invalid quarantine configuration must fail explicitly, never fail open."""
|
||||
|
||||
|
||||
def _manifest():
|
||||
try:
|
||||
raw = QUARANTINE_PATH.read_text(encoding='utf-8')
|
||||
except FileNotFoundError:
|
||||
return {}
|
||||
except (OSError, UnicodeError) as exc:
|
||||
raise QuarantineManifestError('知识库测试索引隔离配置无法读取，请检查 rag-quarantine.json') from exc
|
||||
try:
|
||||
def unique_object(pairs):
|
||||
result = {}
|
||||
for key, value in pairs:
|
||||
if key in result:
|
||||
raise ValueError('duplicate key')
|
||||
result[key] = value
|
||||
return result
|
||||
data = json.loads(raw, object_pairs_hook=unique_object)
|
||||
if not isinstance(data, dict) or set(data) != {'version', 'libraries'}:
|
||||
raise ValueError('invalid fields')
|
||||
if type(data['version']) is not int or data['version'] != 1 or not isinstance(data['libraries'], dict):
|
||||
raise ValueError('invalid version or libraries')
|
||||
result = {}
|
||||
for kid, entries in data['libraries'].items():
|
||||
if not kid.isascii() or not kid.isdecimal() or int(kid) <= 0 or str(int(kid)) != kid or not isinstance(entries, list):
|
||||
raise ValueError('invalid library')
|
||||
names = set()
|
||||
for entry in entries:
|
||||
if not isinstance(entry, dict) or set(entry) != {'file_name', 'reason'}:
|
||||
raise ValueError('invalid entry fields')
|
||||
name, reason = entry['file_name'], entry['reason']
|
||||
if not isinstance(name, str) or not name.strip() or '\x00' in name or name in names:
|
||||
raise ValueError('invalid or duplicate file name')
|
||||
if not isinstance(reason, str) or not reason.strip():
|
||||
raise ValueError('missing reason')
|
||||
names.add(name)
|
||||
result[int(kid)] = names
|
||||
return result
|
||||
except (ValueError, TypeError, AttributeError, OverflowError) as exc:
|
||||
raise QuarantineManifestError('知识库测试索引隔离配置无效，请检查 rag-quarantine.json 的版本、知识库编号及文件条目') from exc
|
||||
|
||||
|
||||
def file_exclusions(session, kid):
|
||||
"""Exclude only explicitly listed names still absent from KnowledgeFile."""
|
||||
if isinstance(kid, bool) or not isinstance(kid, int) or kid <= 0:
|
||||
raise ValueError('知识库编号无效')
|
||||
names = _manifest().get(kid, set())
|
||||
if not names:
|
||||
return set()
|
||||
registered = session.exec(select(KnowledgeFile.file_name).where(KnowledgeFile.kb_id == kid)).all()
|
||||
return names - set(registered)
|
||||
@@ -0,0 +1,115 @@
|
||||
"""常驻检索进程。标准输出只用于 JSON 协议，初始化和检索诊断写标准错误。"""
|
||||
import sys
|
||||
import json
|
||||
import time
|
||||
from contextlib import redirect_stdout
|
||||
from pathlib import Path
|
||||
|
||||
# 从任意 cwd 启动都能找到 app 包（uvicorn 工作目录不保证是 server/）
|
||||
_SERVER_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(_SERVER_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(_SERVER_ROOT))
|
||||
|
||||
|
||||
def initialize():
|
||||
"""每个 worker 仅初始化一次，避免每次查询重复执行迁移和历史表扫描。"""
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError:
|
||||
fcntl = None
|
||||
from app import db as db_mod
|
||||
from app.config import DATA_DIR
|
||||
# 多个首次查询同时启动进程时，不并发执行相同的数据库迁移。
|
||||
with (DATA_DIR / 'rag-worker-init.lock').open('a') as lock:
|
||||
if fcntl:
|
||||
fcntl.flock(lock, fcntl.LOCK_EX)
|
||||
db_mod.init_db()
|
||||
|
||||
|
||||
def handle(payload: dict) -> list:
|
||||
from sqlmodel import Session
|
||||
from app import db as db_mod
|
||||
from app.rag import pipeline as rag_pipeline
|
||||
|
||||
with Session(db_mod.engine) as session:
|
||||
results = rag_pipeline.search_kbs(
|
||||
session,
|
||||
payload["kb_ids"],
|
||||
payload["query"],
|
||||
payload.get("k", 6),
|
||||
excluded_chunk_ids={int(k): set(v) for k, v in (payload.get("excluded") or {}).items()},
|
||||
campaign_id=payload.get("campaign_id", 0),
|
||||
trace_id=payload.get("trace_id", ""),
|
||||
user_id=payload.get("user_id", 0),
|
||||
organization_id=payload.get("organization_id", 1),
|
||||
file_whitelists={int(k): set(v) for k, v in (payload.get("file_whitelists") or {}).items()},
|
||||
)
|
||||
# 透传库名、召回来源、重排分数和各阶段耗时，不裁掉检索诊断信息。
|
||||
return results
|
||||
|
||||
|
||||
def handle_warmup() -> dict:
|
||||
"""Warm the resident worker without issuing a billable retrieval."""
|
||||
from app.rag import embed, store, rerank
|
||||
started = time.monotonic()
|
||||
embed.prewarm()
|
||||
# Optional local Cross-Encoder is loaded in the resident worker so the
|
||||
# first real retrieval does not pay model initialization latency.
|
||||
rerank.prewarm()
|
||||
# Construct the persistent vector client once. Collections remain opened
|
||||
# lazily per KB, while the client and embedding runtime stay resident.
|
||||
store.get_client()
|
||||
return {"warm_state": "warm", "warmup_ms": round((time.monotonic() - started) * 1000, 2)}
|
||||
|
||||
|
||||
def main():
|
||||
protocol_out = sys.stdout
|
||||
try:
|
||||
with redirect_stdout(sys.stderr):
|
||||
initialize()
|
||||
except Exception as exc:
|
||||
from app.system_log import redact
|
||||
protocol_out.write(json.dumps({'ok': False, 'error': '检索进程初始化失败：' + redact(exc, 400)}, ensure_ascii=False) + '\n')
|
||||
protocol_out.flush()
|
||||
return
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
started = time.monotonic()
|
||||
with redirect_stdout(sys.stderr):
|
||||
payload = json.loads(line)
|
||||
if payload.get("op") == "warmup":
|
||||
warm = handle_warmup()
|
||||
protocol_out.write(json.dumps({"ok": True, **warm}, ensure_ascii=False) + "\n")
|
||||
protocol_out.flush()
|
||||
continue
|
||||
out = handle(payload)
|
||||
from app.rag import pipeline as _pipeline
|
||||
metric = _pipeline.last_retrieval_metric()
|
||||
envelope = {"ok": True, "results": out,
|
||||
"worker_timings": {"search_ms": round((time.monotonic() - started) * 1000, 2)},
|
||||
"retrieval_metric": metric}
|
||||
serialization_started = time.perf_counter()
|
||||
json.dumps(envelope, ensure_ascii=False)
|
||||
metric["serialization_ms"] = round((time.perf_counter() - serialization_started) * 1000, 2)
|
||||
envelope["retrieval_metric"] = metric
|
||||
protocol_out.write(json.dumps(envelope, ensure_ascii=False) + "\n")
|
||||
except Exception as exc:
|
||||
from app.system_log import redact
|
||||
metric = {}
|
||||
try:
|
||||
from app.rag import pipeline as _pipeline
|
||||
metric = _pipeline.last_retrieval_metric()
|
||||
except Exception:
|
||||
pass
|
||||
if not metric:
|
||||
metric = {"error": "retrieval_failed"}
|
||||
protocol_out.write(json.dumps({"ok": False, "error": redact(exc, 400),
|
||||
"retrieval_metric": metric}, ensure_ascii=False) + "\n")
|
||||
protocol_out.flush()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,166 @@
|
||||
"""Optional local Cross-Encoder candidate reranking.
|
||||
|
||||
The first-stage vector/FTS recall remains the default. A Cross-Encoder is
|
||||
only enabled when ``RAG_RERANKER_MODEL`` points at a local model directory (or
|
||||
``RAG_RERANKER_ALLOW_DOWNLOAD=1`` is explicitly set). This keeps production
|
||||
search deterministic and offline when the optional model is absent, while
|
||||
allowing operators to add a Chinese model such as ``BAAI/bge-reranker-base``
|
||||
without changing the retrieval contract.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
import threading
|
||||
import time
|
||||
|
||||
logger = logging.getLogger("yawen.rag.rerank")
|
||||
|
||||
_lock = threading.Lock()
|
||||
_metric_lock = threading.Lock()
|
||||
_model = None
|
||||
_model_name = ""
|
||||
_model_error = ""
|
||||
_last_metric: dict = {
|
||||
"enabled": False,
|
||||
"model": "",
|
||||
"model_loaded": False,
|
||||
"model_load_ms": 0.0,
|
||||
"inference_ms": 0.0,
|
||||
"candidate_count": 0,
|
||||
"reranked_count": 0,
|
||||
"fallback": True,
|
||||
}
|
||||
|
||||
|
||||
def _configured_model() -> str:
|
||||
return str(os.environ.get("RAG_RERANKER_MODEL", "")).strip()
|
||||
|
||||
|
||||
def _allow_download() -> bool:
|
||||
return str(os.environ.get("RAG_RERANKER_ALLOW_DOWNLOAD", "0")).lower() in {
|
||||
"1", "true", "yes", "on"
|
||||
}
|
||||
|
||||
|
||||
def _set_metric(**values) -> None:
|
||||
global _last_metric
|
||||
with _metric_lock:
|
||||
row = dict(_last_metric)
|
||||
row.update(values)
|
||||
for key in ("model_load_ms", "inference_ms", "total_ms"):
|
||||
if key in row:
|
||||
try:
|
||||
row[key] = round(float(row[key]), 2)
|
||||
except (TypeError, ValueError):
|
||||
row[key] = 0.0
|
||||
_last_metric = row
|
||||
|
||||
|
||||
def last_metric() -> dict:
|
||||
with _metric_lock:
|
||||
return dict(_last_metric)
|
||||
|
||||
|
||||
def _local_model_path(name: str) -> str | None:
|
||||
"""Resolve a configured model without triggering a network download."""
|
||||
candidate = Path(name).expanduser()
|
||||
if candidate.is_dir() and (candidate / "config.json").exists():
|
||||
return str(candidate)
|
||||
try:
|
||||
from ..config import MODELS_DIR
|
||||
cached = MODELS_DIR / name.replace("/", "--")
|
||||
if cached.is_dir() and (cached / "config.json").exists():
|
||||
return str(cached)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def _get_model():
|
||||
"""Lazily load a local sentence-transformers CrossEncoder."""
|
||||
global _model, _model_name, _model_error
|
||||
name = _configured_model()
|
||||
if not name:
|
||||
return None
|
||||
with _lock:
|
||||
if _model is not None and _model_name == name:
|
||||
return _model
|
||||
if _model_error and _model_name == name:
|
||||
return None
|
||||
started = time.perf_counter()
|
||||
local = _local_model_path(name)
|
||||
if local is None and not _allow_download():
|
||||
_model_name = name
|
||||
_model_error = "模型目录不存在；为避免请求线程下载，已回退基础排序"
|
||||
_set_metric(enabled=False, model=name, model_loaded=False,
|
||||
fallback=True, error=_model_error)
|
||||
logger.warning("RAG 重排模型不可用：%s（%s）", name, _model_error)
|
||||
return None
|
||||
try:
|
||||
from sentence_transformers import CrossEncoder
|
||||
_model = CrossEncoder(local or name)
|
||||
_model_name = name
|
||||
_model_error = ""
|
||||
_set_metric(enabled=True, model=name, model_loaded=True,
|
||||
model_load_ms=(time.perf_counter() - started) * 1000,
|
||||
fallback=False)
|
||||
return _model
|
||||
except Exception as exc: # optional feature must never break RAG
|
||||
_model = None
|
||||
_model_name = name
|
||||
_model_error = f"{type(exc).__name__}: {exc}"
|
||||
_set_metric(enabled=False, model=name, model_loaded=False,
|
||||
model_load_ms=(time.perf_counter() - started) * 1000,
|
||||
fallback=True, error=_model_error)
|
||||
logger.exception("RAG Cross-Encoder 加载失败，回退基础排序：%s", name)
|
||||
return None
|
||||
|
||||
|
||||
def rerank(candidates: list[dict], query: str, *, limit: int | None = None) -> list[dict]:
|
||||
"""Rerank candidates in place-compatible copies and return stable results.
|
||||
|
||||
Only the first ``limit`` candidates are scored; the tail keeps its fused
|
||||
order. If the model is not configured/available, candidates are returned
|
||||
unchanged and no exception escapes the retrieval path.
|
||||
"""
|
||||
configured = _configured_model()
|
||||
if not candidates or not configured:
|
||||
_set_metric(enabled=False, model=configured, candidate_count=len(candidates),
|
||||
reranked_count=0, fallback=True, inference_ms=0.0)
|
||||
return candidates
|
||||
|
||||
count = min(len(candidates), max(1, min(int(limit or os.environ.get("RAG_RERANKER_TOP_K", 32)), 64)))
|
||||
model = _get_model()
|
||||
if model is None:
|
||||
_set_metric(candidate_count=len(candidates), reranked_count=0, fallback=True)
|
||||
return candidates
|
||||
pairs = [(str(query), str(row.get("text") or "")) for row in candidates[:count]]
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
scores = model.predict(pairs, batch_size=min(32, count), show_progress_bar=False)
|
||||
scored = []
|
||||
for index, (row, score) in enumerate(zip(candidates[:count], scores)):
|
||||
item = dict(row)
|
||||
item["rerank_score"] = float(score)
|
||||
scored.append((item, float(score), index))
|
||||
scored.sort(key=lambda value: (-value[1], value[2]))
|
||||
result = [item for item, _, _ in scored] + [dict(row) for row in candidates[count:]]
|
||||
_set_metric(candidate_count=len(candidates), reranked_count=count,
|
||||
inference_ms=(time.perf_counter() - started) * 1000,
|
||||
fallback=False)
|
||||
return result
|
||||
except Exception as exc:
|
||||
logger.exception("RAG Cross-Encoder 推理失败，回退基础排序")
|
||||
_set_metric(candidate_count=len(candidates), reranked_count=0,
|
||||
inference_ms=(time.perf_counter() - started) * 1000,
|
||||
fallback=True, error=f"{type(exc).__name__}: {exc}")
|
||||
return candidates
|
||||
|
||||
|
||||
def prewarm() -> bool:
|
||||
"""Load the optional reranker in the resident RAG worker at startup."""
|
||||
if not _configured_model():
|
||||
return False
|
||||
return _get_model() is not None
|
||||
@@ -0,0 +1,788 @@
|
||||
"""RAG 向量索引（collection=知识库），显式选择 Chroma 或 Qdrant。"""
|
||||
import hashlib
|
||||
import logging
|
||||
import re
|
||||
import shutil
|
||||
import sqlite3
|
||||
import json
|
||||
import math
|
||||
import time
|
||||
import uuid
|
||||
import threading
|
||||
from ..config import CHROMA_DIR, DATA_DIR
|
||||
from . import backend
|
||||
|
||||
logger = logging.getLogger("yawen.rag.store")
|
||||
|
||||
|
||||
class VectorIndexDisabled(RuntimeError):
|
||||
"""Raised when an operation requires a vector backend but it is disabled."""
|
||||
|
||||
|
||||
_client = None
|
||||
_client_init_lock = threading.Lock()
|
||||
FALLBACK_DIR = DATA_DIR / "vector_fallback"
|
||||
FALLBACK_DIR.mkdir(parents=True, exist_ok=True)
|
||||
_count_cache: dict[int, tuple[str, dict[str, int]]] = {}
|
||||
_search_metric_local = threading.local()
|
||||
COUNT_CACHE_DIR = DATA_DIR / "vector_counts"
|
||||
COUNT_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _count_cache_path(kb_id: int):
|
||||
return COUNT_CACHE_DIR / f"kb_{int(kb_id)}.json"
|
||||
|
||||
def kb_revision(kb_id: int) -> str:
|
||||
"""Cross-process version of one library, namespaced to this data directory."""
|
||||
marker = COUNT_CACHE_DIR / f"kb_{int(kb_id)}.revision"
|
||||
parts = [backend.namespace(), str(CHROMA_DIR.resolve()), str(FALLBACK_DIR.resolve()), str(marker.resolve())]
|
||||
try:
|
||||
parts.append(marker.read_text(encoding="ascii"))
|
||||
except FileNotFoundError:
|
||||
parts.append("initial")
|
||||
if not backend.is_qdrant():
|
||||
fallback = _fallback_path(kb_id)
|
||||
try:
|
||||
stat = fallback.stat()
|
||||
parts.append(f"{stat.st_ino}:{stat.st_size}:{stat.st_mtime_ns}:{stat.st_ctime_ns}")
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
return hashlib.sha256("\0".join(parts).encode()).hexdigest()
|
||||
|
||||
|
||||
def _bump_revision(kb_id: int) -> None:
|
||||
COUNT_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
marker = COUNT_CACHE_DIR / f"kb_{int(kb_id)}.revision"
|
||||
temporary = marker.with_suffix(f".{uuid.uuid4().hex}.tmp")
|
||||
try:
|
||||
temporary.write_text(uuid.uuid4().hex, encoding="ascii")
|
||||
temporary.replace(marker)
|
||||
finally:
|
||||
temporary.unlink(missing_ok=True)
|
||||
|
||||
|
||||
def _invalidate_count_cache(kb_id: int) -> None:
|
||||
_bump_revision(kb_id)
|
||||
_count_cache.pop(int(kb_id), None)
|
||||
try:
|
||||
_count_cache_path(kb_id).unlink(missing_ok=True)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _fallback_path(kb_id: int):
|
||||
return FALLBACK_DIR / f"kb_{int(kb_id)}.json"
|
||||
|
||||
|
||||
def _read_fallback(kb_id: int) -> dict | None:
|
||||
if backend.is_qdrant():
|
||||
return None
|
||||
path = _fallback_path(kb_id)
|
||||
if not path.exists():
|
||||
return None
|
||||
try:
|
||||
payload = json.loads(path.read_text(encoding="utf-8"))
|
||||
if isinstance(payload, dict) and isinstance(payload.get("items"), list):
|
||||
return payload
|
||||
except (OSError, ValueError, TypeError):
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _write_fallback(kb_id: int, items: list[dict]) -> None:
|
||||
path = _fallback_path(kb_id)
|
||||
temp = path.with_suffix(".tmp")
|
||||
payload = {"version": 1, "kb_id": int(kb_id), "items": items}
|
||||
temp.write_text(json.dumps(payload, ensure_ascii=False, separators=(",", ":")), encoding="utf-8")
|
||||
temp.replace(path)
|
||||
_bump_revision(kb_id)
|
||||
|
||||
|
||||
class _StreamingJSON:
|
||||
"""只保留当前字段的 JSON 流；向量字段逐字符跳过，不构造 embedding 数组。"""
|
||||
def __init__(self, handle):
|
||||
self.handle, self.buffer, self.pos, self.eof = handle, "", 0, False
|
||||
|
||||
def peek(self):
|
||||
if self.pos >= len(self.buffer) and not self.eof:
|
||||
self.buffer, self.pos = self.handle.read(65536), 0
|
||||
self.eof = not self.buffer
|
||||
return self.buffer[self.pos:self.pos + 1]
|
||||
|
||||
def take(self):
|
||||
value = self.peek()
|
||||
self.pos += bool(value)
|
||||
return value
|
||||
|
||||
def whitespace(self):
|
||||
while self.peek() and self.peek().isspace():
|
||||
self.take()
|
||||
|
||||
def expect(self, value):
|
||||
self.whitespace()
|
||||
if self.take() != value:
|
||||
raise ValueError("索引 JSON 结构损坏")
|
||||
|
||||
def value(self, retain=True):
|
||||
self.whitespace()
|
||||
first = self.peek()
|
||||
if not first:
|
||||
raise ValueError("索引 JSON 意外结束")
|
||||
output, depth, quoted, escaped = [], 0, False, False
|
||||
while self.peek():
|
||||
char = self.peek()
|
||||
if not quoted and depth == 0 and output and (char in ",]}" or char.isspace()):
|
||||
break
|
||||
# 不保留的值也需要一个已开始标记，避免把后续逗号吞掉。
|
||||
if not quoted and depth == 0 and not retain and output and (char in ",]}" or char.isspace()):
|
||||
break
|
||||
char = self.take()
|
||||
if retain:
|
||||
output.append(char)
|
||||
elif not output:
|
||||
output.append(True)
|
||||
if quoted:
|
||||
if escaped:
|
||||
escaped = False
|
||||
elif char == "\\":
|
||||
escaped = True
|
||||
elif char == '"':
|
||||
quoted = False
|
||||
elif char == '"':
|
||||
quoted = True
|
||||
elif char in "[{":
|
||||
depth += 1
|
||||
elif char in "]}":
|
||||
depth -= 1
|
||||
if depth == 0 and not quoted and first in '[{"':
|
||||
break
|
||||
if quoted or depth:
|
||||
raise ValueError("索引 JSON 字段不完整")
|
||||
return json.loads("".join(output)) if retain else None
|
||||
|
||||
|
||||
def _stream_fallback(kb_id: int):
|
||||
path = _fallback_path(kb_id)
|
||||
before = path.stat()
|
||||
with path.open(encoding="utf-8") as handle:
|
||||
stream, found = _StreamingJSON(handle), False
|
||||
stream.expect("{")
|
||||
while True:
|
||||
stream.whitespace()
|
||||
if stream.peek() == "}":
|
||||
stream.take()
|
||||
break
|
||||
key = stream.value()
|
||||
stream.expect(":")
|
||||
if key == "items":
|
||||
if found:
|
||||
raise ValueError("索引包含重复 items 字段")
|
||||
found = True
|
||||
stream.expect("[")
|
||||
while True:
|
||||
stream.whitespace()
|
||||
if stream.peek() == "]":
|
||||
stream.take()
|
||||
break
|
||||
stream.expect("{")
|
||||
item = {}
|
||||
while True:
|
||||
stream.whitespace()
|
||||
if stream.peek() == "}":
|
||||
stream.take()
|
||||
break
|
||||
field = stream.value()
|
||||
stream.expect(":")
|
||||
if field in {"id", "text", "meta"}:
|
||||
item[field] = stream.value()
|
||||
else:
|
||||
stream.value(retain=False)
|
||||
stream.whitespace()
|
||||
if stream.peek() == "}":
|
||||
continue
|
||||
stream.expect(",")
|
||||
yield item
|
||||
stream.whitespace()
|
||||
if stream.peek() == "]":
|
||||
continue
|
||||
stream.expect(",")
|
||||
elif key == "kb_id":
|
||||
if stream.value() != kb_id:
|
||||
raise ValueError("索引知识库编号不匹配")
|
||||
else:
|
||||
stream.value(retain=False)
|
||||
stream.whitespace()
|
||||
if stream.peek() == "}":
|
||||
continue
|
||||
stream.expect(",")
|
||||
stream.whitespace()
|
||||
if stream.peek() or not found:
|
||||
raise ValueError("索引 JSON 缺少正文或包含多余数据")
|
||||
after = path.stat()
|
||||
if (before.st_ino, before.st_size, before.st_mtime_ns) != (after.st_ino, after.st_size, after.st_mtime_ns):
|
||||
raise ValueError("读取过程中知识库索引发生变化，请在索引完成后重新评估")
|
||||
|
||||
|
||||
def iter_chunks(kb_id: int, page_size: int = 256, file_names: set[str] | None = None):
|
||||
"""顺序遍历完整索引正文；不做相似度筛选，不加载向量，不改变消费记录。"""
|
||||
if backend.is_disabled():
|
||||
return
|
||||
if isinstance(kb_id, bool) or not isinstance(kb_id, int) or kb_id <= 0:
|
||||
raise ValueError("知识库编号无效")
|
||||
if not isinstance(page_size, int) or not 1 <= page_size <= 4000:
|
||||
raise ValueError("索引分页大小无效")
|
||||
if file_names is not None and not file_names:
|
||||
return
|
||||
if backend.is_qdrant():
|
||||
revision = kb_revision(kb_id)
|
||||
collection = get_client().get_collection(collection_name(kb_id))
|
||||
total = collection.count()
|
||||
seen = set()
|
||||
groups = [None] if file_names is None else [sorted(file_names)[i:i + 128] for i in range(0, len(file_names), 128)]
|
||||
for names in groups:
|
||||
where = None if names is None else {"file": {"$in": names}}
|
||||
for page in collection.iter_records(where=where, page_size=page_size):
|
||||
for index, identifier in enumerate(page['ids']):
|
||||
meta = page['metadatas'][index]
|
||||
if identifier in seen:
|
||||
raise RuntimeError("知识索引分页返回重复片段，请重新评估")
|
||||
seen.add(identifier)
|
||||
if meta.get('kb_id') not in (None, kb_id):
|
||||
raise RuntimeError("知识索引返回了其他知识库的片段")
|
||||
if file_names is not None and meta.get('file') not in file_names:
|
||||
raise RuntimeError("知识索引返回了分类范围外的文件")
|
||||
yield {"id": identifier, "text": page['documents'][index], "meta": meta}
|
||||
if collection.count() != total or kb_revision(kb_id) != revision or (file_names is None and len(seen) != total):
|
||||
raise RuntimeError("读取过程中知识索引数量变化，请重新评估")
|
||||
return
|
||||
if _fallback_path(kb_id).exists():
|
||||
try:
|
||||
for item in _stream_fallback(kb_id):
|
||||
if file_names is None or (item.get('meta') or {}).get('file') in file_names:
|
||||
yield item
|
||||
except (OSError, ValueError, TypeError) as exc:
|
||||
raise RuntimeError(f"知识库 {kb_id} 本地索引读取失败：{exc}") from exc
|
||||
return
|
||||
try:
|
||||
collection = get_client().get_collection(collection_name(kb_id))
|
||||
total, offset = collection.count(), 0
|
||||
if file_names is not None:
|
||||
# Filter in Chroma before paging; only selected-category documents are read.
|
||||
names = sorted(file_names)
|
||||
for start in range(0, len(names), 128):
|
||||
offset = 0
|
||||
while True:
|
||||
page = collection.get(include=["documents", "metadatas"], limit=page_size, offset=offset,
|
||||
where={"file": {"$in": names[start:start + 128]}})
|
||||
ids = _as_list(page.get("ids"))
|
||||
documents, metadata = _as_list(page.get("documents")), _as_list(page.get("metadatas"))
|
||||
if len(documents) != len(ids) or len(metadata) != len(ids):
|
||||
raise ValueError("分类索引正文或来源缺失")
|
||||
for index, identifier in enumerate(ids):
|
||||
meta = metadata[index] or {}
|
||||
if meta.get('file') not in file_names:
|
||||
raise ValueError("分类索引返回了范围外的文件")
|
||||
yield {"id": identifier, "text": documents[index], "meta": meta}
|
||||
offset += len(ids)
|
||||
if len(ids) < page_size:
|
||||
break
|
||||
if collection.count() != total:
|
||||
raise ValueError("读取过程中索引数量变化，请重新评估")
|
||||
return
|
||||
while offset < total:
|
||||
page = collection.get(include=["documents", "metadatas"], limit=page_size, offset=offset)
|
||||
ids = _as_list(page.get("ids"))
|
||||
documents, metadata = _as_list(page.get("documents")), _as_list(page.get("metadatas"))
|
||||
if not ids or len(documents) != len(ids) or len(metadata) != len(ids):
|
||||
raise ValueError("索引分页正文或来源缺失")
|
||||
for index, identifier in enumerate(ids):
|
||||
yield {"id": identifier, "text": documents[index], "meta": metadata[index] or {}}
|
||||
offset += len(ids)
|
||||
if offset != total or collection.count() != total:
|
||||
raise ValueError("读取过程中索引数量变化，请重新评估")
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"知识库 {kb_id} 索引读取失败：{exc}") from exc
|
||||
|
||||
def _as_list(value, default: list | None = None) -> list:
|
||||
"""把 Chroma/NumPy 返回的序列安全转换为普通列表，避免 array 的布尔判断异常。"""
|
||||
if value is None:
|
||||
return list(default or [])
|
||||
if isinstance(value, list):
|
||||
return value
|
||||
try:
|
||||
return list(value)
|
||||
except TypeError:
|
||||
return list(default or [])
|
||||
|
||||
def get_client():
|
||||
if backend.is_disabled():
|
||||
configured = backend.configured_backend()
|
||||
if configured == "chroma":
|
||||
raise VectorIndexDisabled("本地向量索引未启用：未安装 chromadb，请使用轻量模式或安装可选 RAG 依赖")
|
||||
raise VectorIndexDisabled("本地向量索引已禁用：设置 RAG_VECTOR_BACKEND=chroma/qdrant 后才可访问向量库")
|
||||
if backend.is_qdrant():
|
||||
from .qdrant_store import get_client as remote_client
|
||||
return remote_client()
|
||||
global _client
|
||||
if _client is None:
|
||||
with _client_init_lock:
|
||||
if _client is None:
|
||||
try:
|
||||
import chromadb
|
||||
except ImportError as exc:
|
||||
raise VectorIndexDisabled("缺少 chromadb，本地向量索引不可用") from exc
|
||||
_client = chromadb.PersistentClient(path=str(CHROMA_DIR))
|
||||
return _client
|
||||
|
||||
def collection_name(kb_id: int) -> str:
|
||||
return f"kb_{kb_id}"
|
||||
|
||||
def ensure_collection(kb_id: int):
|
||||
started = time.perf_counter()
|
||||
collection = get_client().get_or_create_collection(collection_name(kb_id))
|
||||
metric = getattr(_search_metric_local, "metric", None)
|
||||
if metric is not None:
|
||||
metric["collection_get_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||||
return collection
|
||||
|
||||
|
||||
def last_search_metric() -> dict:
|
||||
"""Return timings for the most recent vector search on this thread.
|
||||
|
||||
Search requests fan out across threads, so a thread-local envelope avoids
|
||||
cross-library races while keeping the public ``search`` signature intact.
|
||||
"""
|
||||
return dict(getattr(_search_metric_local, "metric", {}) or {})
|
||||
|
||||
def _chunk_id(file_name: str, item: dict, occurrence: int) -> str:
|
||||
"""用来源和正文生成稳定 ID，重建整个库也不会改变同一片段的消费记录。"""
|
||||
meta = item.get("meta", {})
|
||||
heading = str(meta.get("heading") or "")
|
||||
text = " ".join(str(item.get("text") or "").split())
|
||||
raw = f"{file_name}\0{heading}\0{text}\0{occurrence}".encode("utf-8")
|
||||
return "chunk-" + hashlib.sha256(raw).hexdigest()
|
||||
|
||||
def _file_rows(collection, file_name: str) -> list[dict]:
|
||||
if hasattr(collection, "iter_records"):
|
||||
total = collection.count(where={"file": file_name})
|
||||
rows = []
|
||||
for data in collection.iter_records(where={"file": file_name}, include=["documents", "metadatas", "embeddings"]):
|
||||
rows.extend({"id": identifier, "text": data["documents"][i], "embedding": data["embeddings"][i], "meta": data["metadatas"][i]}
|
||||
for i, identifier in enumerate(data["ids"]))
|
||||
if len(rows) != total or collection.count(where={"file": file_name}) != total:
|
||||
raise RuntimeError("文件索引快照不完整，已取消修改")
|
||||
return rows
|
||||
rows, offset = [], 0
|
||||
while True:
|
||||
data = collection.get(where={"file": file_name}, include=["documents", "embeddings", "metadatas"],
|
||||
limit=4000, offset=offset)
|
||||
ids = _as_list(data.get("ids"))
|
||||
texts, vectors, metas = (_as_list(data.get(key)) for key in ("documents", "embeddings", "metadatas"))
|
||||
if any(len(values) != len(ids) for values in (texts, vectors, metas)):
|
||||
raise RuntimeError("文件索引快照不完整，已取消修改")
|
||||
rows.extend({"id": ident, "text": texts[i], "embedding": _as_list(vectors[i]), "meta": metas[i] or {}}
|
||||
for i, ident in enumerate(ids))
|
||||
offset += len(ids)
|
||||
if len(ids) < 4000:
|
||||
return rows
|
||||
|
||||
|
||||
def snapshot_file(kb_id: int, file_name: str) -> dict:
|
||||
"""Snapshot the file in each existing index; never seed a partial fallback."""
|
||||
if backend.is_disabled():
|
||||
return {"kb_id": kb_id, "file_name": file_name, "fallback_exists": False,
|
||||
"backend_namespace": backend.namespace(), "fallback_items": [], "chroma_items": []}
|
||||
exists = not backend.is_qdrant() and _fallback_path(kb_id).exists()
|
||||
fallback = _read_fallback(kb_id) if exists else None
|
||||
if exists and fallback is None:
|
||||
raise RuntimeError("本地向量索引损坏，已取消修改")
|
||||
return {"kb_id": kb_id, "file_name": file_name, "fallback_exists": exists,
|
||||
"backend_namespace": backend.namespace(),
|
||||
"fallback_items": [row for row in (fallback or {}).get("items", [])
|
||||
if (row.get("meta") or {}).get("file") == file_name],
|
||||
"chroma_items": _file_rows(ensure_collection(kb_id), file_name)}
|
||||
|
||||
|
||||
def _insert_rows(collection, rows: list[dict]):
|
||||
for start in range(0, len(rows), 4000):
|
||||
batch = rows[start:start + 4000]
|
||||
collection.add(ids=[row["id"] for row in batch], documents=[row["text"] for row in batch],
|
||||
embeddings=[row["embedding"] for row in batch], metadatas=[row["meta"] for row in batch])
|
||||
|
||||
|
||||
def restore_file(kb_id: int, file_name: str, snapshot: dict) -> None:
|
||||
if snapshot.get("kb_id") != kb_id or snapshot.get("file_name") != file_name:
|
||||
raise ValueError("文件索引快照归属不匹配")
|
||||
if snapshot.get("backend_namespace", backend.namespace()) != backend.namespace():
|
||||
raise ValueError("文件索引快照所属后端已改变，已取消恢复")
|
||||
collection = ensure_collection(kb_id)
|
||||
try:
|
||||
collection.delete(where={"file": file_name})
|
||||
_insert_rows(collection, snapshot["chroma_items"])
|
||||
if snapshot["fallback_exists"]:
|
||||
fallback = _read_fallback(kb_id)
|
||||
if fallback is None:
|
||||
raise RuntimeError("回滚时本地索引不可读取")
|
||||
others = [row for row in fallback["items"] if (row.get("meta") or {}).get("file") != file_name]
|
||||
_write_fallback(kb_id, others + snapshot["fallback_items"])
|
||||
finally:
|
||||
_invalidate_count_cache(kb_id)
|
||||
|
||||
|
||||
def _prepared_rows(kb_id, chunks, embeddings, file_name, category_path):
|
||||
if len(chunks) != len(embeddings):
|
||||
raise ValueError("分段数量与向量数量不一致，已取消写入")
|
||||
rows, occurrences = [], {}
|
||||
for i, item in enumerate(chunks):
|
||||
if not isinstance(item.get("text"), str) or not item["text"].strip():
|
||||
raise ValueError("索引片段正文为空，已取消写入")
|
||||
vector = _as_list(embeddings[i])
|
||||
if not vector or any(not math.isfinite(float(x)) for x in vector):
|
||||
raise ValueError("索引向量无效，已取消写入")
|
||||
if backend.is_qdrant() and len(vector) != 512:
|
||||
raise ValueError("Qdrant 索引向量必须是 512 维，已取消写入")
|
||||
fingerprint = f"{item.get('meta', {}).get('heading', '')}\0{' '.join(item['text'].split())}"
|
||||
occurrences[fingerprint] = occurrences.get(fingerprint, 0) + 1
|
||||
ident = _chunk_id(file_name, item, occurrences[fingerprint])
|
||||
meta = dict(item.get("meta") or {})
|
||||
meta.setdefault("seq", i)
|
||||
from .embed import EMBEDDING_POLICY
|
||||
meta.update(kb_id=kb_id, file=file_name, chunk_hash=ident.removeprefix("chunk-"), embedding_policy=EMBEDDING_POLICY)
|
||||
if category_path:
|
||||
meta["category_path"] = category_path
|
||||
rows.append({"id": ident, "text": item["text"], "embedding": vector, "meta": meta})
|
||||
return rows
|
||||
|
||||
|
||||
def replace_file_chunks(kb_id: int, chunks: list[dict], embeddings: list[list[float]],
|
||||
file_name: str, category_path: str = "") -> None:
|
||||
if backend.is_disabled():
|
||||
raise VectorIndexDisabled("本地向量索引已禁用，不能写入向量片段")
|
||||
"""Prepare before mutation; failed replacements restore that file in both stores.
|
||||
|
||||
Chroma-only libraries stay Chroma-only. A fallback, when already present,
|
||||
continues to contain every other file rather than shadowing the whole library
|
||||
with only this upload. Callers serialize writers with the per-KB import lock.
|
||||
"""
|
||||
rows = _prepared_rows(kb_id, chunks, embeddings, file_name, category_path)
|
||||
previous = snapshot_file(kb_id, file_name)
|
||||
collection = ensure_collection(kb_id)
|
||||
_invalidate_count_cache(kb_id)
|
||||
try:
|
||||
collection.delete(where={"file": file_name})
|
||||
_insert_rows(collection, rows)
|
||||
if previous["fallback_exists"]:
|
||||
fallback = _read_fallback(kb_id)
|
||||
if fallback is None:
|
||||
raise RuntimeError("本地向量索引无法读取")
|
||||
others = [row for row in fallback["items"] if (row.get("meta") or {}).get("file") != file_name]
|
||||
_write_fallback(kb_id, others + rows)
|
||||
_invalidate_count_cache(kb_id)
|
||||
except Exception as exc:
|
||||
try:
|
||||
restore_file(kb_id, file_name, previous)
|
||||
except Exception as recovery:
|
||||
raise RuntimeError(f"文件索引更新失败，恢复也失败，请保留源文件并检查索引：{recovery}") from exc
|
||||
raise RuntimeError(f"文件索引更新失败，已恢复原索引：{exc}") from exc
|
||||
|
||||
|
||||
def add_chunks(kb_id: int, chunks: list[dict], embeddings: list[list[float]],
|
||||
file_name: str, category_path: str = ""):
|
||||
return replace_file_chunks(kb_id, chunks, embeddings, file_name, category_path)
|
||||
|
||||
|
||||
def delete_file(kb_id: int, file_name: str):
|
||||
if backend.is_disabled():
|
||||
_invalidate_count_cache(kb_id)
|
||||
return
|
||||
return replace_file_chunks(kb_id, [], [], file_name)
|
||||
|
||||
def count_by_file(kb_id: int) -> dict:
|
||||
"""按当前后端及版本缓存文件片段数；Qdrant 故障不读取旧 Chroma 缓存。"""
|
||||
if backend.is_disabled():
|
||||
return {}
|
||||
remote = backend.is_qdrant()
|
||||
revision = kb_revision(kb_id)
|
||||
cache_path = _count_cache_path(kb_id)
|
||||
try:
|
||||
# An exact remote count also verifies service availability before cached
|
||||
# statistics are returned. The endpoint is cheap and does not load text.
|
||||
coll = ensure_collection(kb_id) if remote else None
|
||||
total = coll.count() if remote else None
|
||||
cached = _count_cache.get(int(kb_id))
|
||||
if cached and cached[0] == revision and (total is None or sum(cached[1].values()) == total):
|
||||
return dict(cached[1])
|
||||
try:
|
||||
payload = json.loads(cache_path.read_text(encoding="utf-8"))
|
||||
counts = payload.get("counts") if isinstance(payload, dict) else None
|
||||
if payload.get("version") == 2 and payload.get("kb_id") == int(kb_id) and payload.get("revision") == revision and isinstance(counts, dict):
|
||||
if any(not isinstance(k, str) or isinstance(v, bool) or not isinstance(v, int) or v < 0 for k, v in counts.items()):
|
||||
raise ValueError("知识索引统计缓存无效")
|
||||
if total is None or sum(counts.values()) == total:
|
||||
_count_cache[int(kb_id)] = (revision, counts)
|
||||
return dict(counts)
|
||||
except (OSError, ValueError, TypeError, AttributeError):
|
||||
pass
|
||||
out: dict[str, int] = {}
|
||||
if not remote and _fallback_path(kb_id).exists():
|
||||
records = (item.get("meta") or {} for item in _stream_fallback(kb_id))
|
||||
for meta in records:
|
||||
name = meta.get("file") or "未知"
|
||||
out[name] = out.get(name, 0) + 1
|
||||
else:
|
||||
coll = coll if coll is not None else ensure_collection(kb_id)
|
||||
if coll is None:
|
||||
return {}
|
||||
if hasattr(coll, "iter_records"):
|
||||
pages = coll.iter_records(include=["metadatas"])
|
||||
for data in pages:
|
||||
for meta in data["metadatas"]:
|
||||
name = meta.get("file") or "未知"
|
||||
out[name] = out.get(name, 0) + 1
|
||||
else:
|
||||
offset, page = 0, 5000
|
||||
while True:
|
||||
got = coll.get(include=["metadatas"], limit=page, offset=offset)
|
||||
metas = got.get("metadatas") or []
|
||||
for meta in metas:
|
||||
name = (meta or {}).get("file") or "未知"
|
||||
out[name] = out.get(name, 0) + 1
|
||||
offset += len(metas)
|
||||
if len(metas) < page:
|
||||
break
|
||||
if kb_revision(kb_id) != revision or (remote and (sum(out.values()) != total or coll.count() != total)):
|
||||
raise RuntimeError("统计过程中知识索引发生变化，请重新加载")
|
||||
_count_cache[int(kb_id)] = (revision, out)
|
||||
temporary = cache_path.with_suffix(f".{uuid.uuid4().hex}.tmp")
|
||||
try:
|
||||
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary.write_text(json.dumps({"version": 2, "kb_id": int(kb_id), "revision": revision, "counts": out},
|
||||
ensure_ascii=False, separators=(",", ":")), encoding="utf-8")
|
||||
temporary.replace(cache_path)
|
||||
except OSError:
|
||||
pass
|
||||
finally:
|
||||
temporary.unlink(missing_ok=True)
|
||||
return dict(out)
|
||||
except Exception:
|
||||
if remote:
|
||||
raise
|
||||
return {}
|
||||
|
||||
|
||||
def count(kb_id: int) -> int:
|
||||
if backend.is_disabled():
|
||||
return 0
|
||||
fallback = _read_fallback(kb_id)
|
||||
if fallback is not None:
|
||||
return len(fallback.get("items", []))
|
||||
try:
|
||||
return ensure_collection(kb_id).count()
|
||||
except Exception:
|
||||
if backend.is_qdrant():
|
||||
raise
|
||||
return 0
|
||||
|
||||
def search(kb_id: int, query_embedding: list[float], k: int = 6, file_names: set[str] | None = None) -> list[dict]:
|
||||
if backend.is_disabled():
|
||||
return []
|
||||
metric = {"kb_id": int(kb_id), "metadata_filter_ms": 0.0,
|
||||
"vector_ms": 0.0, "collection_get_ms": 0.0}
|
||||
_search_metric_local.metric = metric
|
||||
started = time.perf_counter()
|
||||
if file_names is not None and not file_names:
|
||||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||||
return []
|
||||
fallback = _read_fallback(kb_id)
|
||||
if fallback is not None:
|
||||
filter_started = time.perf_counter()
|
||||
if file_names is not None:
|
||||
filtered = [item for item in fallback.get("items", [])
|
||||
if (item.get("meta") or {}).get("file") in file_names]
|
||||
else:
|
||||
filtered = list(fallback.get("items", []))
|
||||
metric["metadata_filter_ms"] = round((time.perf_counter() - filter_started) * 1000, 2)
|
||||
vector_started = time.perf_counter()
|
||||
def cosine_distance(vector: list[float]) -> float:
|
||||
dot = sum(float(a) * float(b) for a, b in zip(query_embedding, vector))
|
||||
left = math.sqrt(sum(float(a) * float(a) for a in query_embedding)) or 1.0
|
||||
right = math.sqrt(sum(float(b) * float(b) for b in vector)) or 1.0
|
||||
return 1.0 - dot / (left * right)
|
||||
ranked = sorted(filtered, key=lambda item: cosine_distance(item.get("embedding", [])))[:k]
|
||||
metric["vector_ms"] = round((time.perf_counter() - vector_started) * 1000, 2)
|
||||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||||
return [{"id": item.get("id"), "text": item.get("text", ""),
|
||||
"meta": item.get("meta", {}), "distance": cosine_distance(item.get("embedding", []))}
|
||||
for item in ranked]
|
||||
coll = ensure_collection(kb_id)
|
||||
# Chroma and Qdrant both return an empty result for an empty collection.
|
||||
# Do not call ``count()`` on the hot retrieval path: on an external disk
|
||||
# it can scan collection metadata (over 1.7s cold for the 108k-chunk KB),
|
||||
# while the bounded nearest-neighbor query itself is only a few ms.
|
||||
# Keeping ``available`` unknown also removes a redundant remote count for
|
||||
# Qdrant; the backend safely truncates to the number of stored points.
|
||||
available = None
|
||||
if file_names is not None:
|
||||
candidates = []
|
||||
names = sorted(file_names)
|
||||
for start in range(0, len(names), 128):
|
||||
vector_started = time.perf_counter()
|
||||
result = coll.query(query_embeddings=[query_embedding], n_results=k,
|
||||
where={"file": {"$in": names[start:start + 128]}})
|
||||
metric["vector_ms"] += (time.perf_counter() - vector_started) * 1000
|
||||
filter_started = time.perf_counter()
|
||||
for i, identifier in enumerate((result.get('ids') or [[]])[0]):
|
||||
meta = result['metadatas'][0][i] or {}
|
||||
if meta.get('file') in file_names:
|
||||
candidates.append({'id': identifier, 'text': result['documents'][0][i], 'meta': meta,
|
||||
'distance': result['distances'][0][i]})
|
||||
metric["metadata_filter_ms"] += (time.perf_counter() - filter_started) * 1000
|
||||
metric["metadata_filter_ms"] = round(metric["metadata_filter_ms"], 2)
|
||||
metric["vector_ms"] = round(metric["vector_ms"], 2)
|
||||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||||
return sorted(candidates, key=lambda item: item['distance'])[:k]
|
||||
vector_started = time.perf_counter()
|
||||
res = coll.query(query_embeddings=[query_embedding], n_results=k)
|
||||
metric["vector_ms"] = round((time.perf_counter() - vector_started) * 1000, 2)
|
||||
out = []
|
||||
if res and res.get("ids") and res["ids"][0]:
|
||||
for i, doc_id in enumerate(res["ids"][0]):
|
||||
out.append({
|
||||
"id": doc_id,
|
||||
"text": res["documents"][0][i],
|
||||
"meta": res["metadatas"][0][i] if res.get("metadatas") else {},
|
||||
"distance": res["distances"][0][i] if res.get("distances") else None,
|
||||
})
|
||||
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||||
return out
|
||||
|
||||
def _vector_segment_dirs(kb_id: int) -> list:
|
||||
"""在删除 collection 前记录其 HNSW 段目录；只返回当前 collection 的 UUID 目录。"""
|
||||
db_path = CHROMA_DIR / "chroma.sqlite3"
|
||||
if not db_path.exists():
|
||||
return []
|
||||
with sqlite3.connect(str(db_path)) as conn:
|
||||
rows = conn.execute(
|
||||
"""SELECT s.id FROM collections c JOIN segments s ON s.collection = c.id
|
||||
WHERE c.name = ? AND s.scope = 'VECTOR'""", (collection_name(kb_id),)
|
||||
).fetchall()
|
||||
root = CHROMA_DIR.resolve()
|
||||
dirs = []
|
||||
for (segment_id,) in rows:
|
||||
if not re.fullmatch(r"[0-9a-f-]{36}", str(segment_id)):
|
||||
continue
|
||||
candidate = (root / str(segment_id)).resolve()
|
||||
if candidate.parent == root:
|
||||
dirs.append(candidate)
|
||||
return dirs
|
||||
|
||||
|
||||
def reset_kb(kb_id: int):
|
||||
if backend.is_disabled():
|
||||
_fallback_path(kb_id).unlink(missing_ok=True)
|
||||
_invalidate_count_cache(kb_id)
|
||||
return
|
||||
if backend.is_qdrant():
|
||||
_invalidate_count_cache(kb_id)
|
||||
get_client().delete_collection(collection_name(kb_id))
|
||||
_invalidate_count_cache(kb_id)
|
||||
return
|
||||
fallback_path = _fallback_path(kb_id)
|
||||
if fallback_path.exists():
|
||||
fallback_path.unlink()
|
||||
_invalidate_count_cache(kb_id)
|
||||
segment_dirs = _vector_segment_dirs(kb_id)
|
||||
try:
|
||||
get_client().delete_collection(collection_name(kb_id))
|
||||
except Exception:
|
||||
# collection 不存在时无需清理；其余失败交给调用方，避免误删仍被引用的段目录。
|
||||
if segment_dirs:
|
||||
raise
|
||||
return
|
||||
_invalidate_count_cache(kb_id)
|
||||
for directory in segment_dirs:
|
||||
if directory.exists():
|
||||
shutil.rmtree(directory)
|
||||
|
||||
|
||||
def snapshot_kb(kb_id: int) -> dict:
|
||||
"""重建前导出当前 collection，供分块/嵌入异常时原样恢复。"""
|
||||
if backend.is_disabled():
|
||||
return {"ids": [], "documents": [], "embeddings": [], "metadatas": []}
|
||||
fallback = _read_fallback(kb_id)
|
||||
if fallback is not None:
|
||||
return {"ids": [item.get("id") for item in fallback["items"]],
|
||||
"documents": [item.get("text", "") for item in fallback["items"]],
|
||||
"embeddings": [item.get("embedding", []) for item in fallback["items"]],
|
||||
"metadatas": [item.get("meta", {}) for item in fallback["items"]]}
|
||||
coll = ensure_collection(kb_id)
|
||||
if coll.count() <= 0:
|
||||
return {"ids": [], "documents": [], "embeddings": [], "metadatas": []}
|
||||
if hasattr(coll, "iter_records"):
|
||||
revision, total = kb_revision(kb_id), coll.count()
|
||||
result = {"ids": [], "documents": [], "embeddings": [], "metadatas": []}
|
||||
for page in coll.iter_records(include=["documents", "metadatas", "embeddings"]):
|
||||
for key in result:
|
||||
result[key].extend(page[key])
|
||||
if len(result["ids"]) != total or coll.count() != total or kb_revision(kb_id) != revision:
|
||||
raise RuntimeError("知识库索引快照不完整，已取消重建")
|
||||
return result
|
||||
# 大库分页导出（单次全量 get 会触发 sqlite 变量上限）
|
||||
ids: list = []
|
||||
documents: list = []
|
||||
embeddings: list = []
|
||||
metadatas: list = []
|
||||
page = 4000
|
||||
offset = 0
|
||||
while True:
|
||||
data = coll.get(include=["documents", "embeddings", "metadatas"], limit=page, offset=offset)
|
||||
batch_ids = _as_list(data.get("ids"))
|
||||
if not batch_ids:
|
||||
break
|
||||
ids.extend(batch_ids)
|
||||
documents.extend(_as_list(data.get("documents")))
|
||||
embeddings.extend(_as_list(data.get("embeddings")))
|
||||
metadatas.extend(_as_list(data.get("metadatas")))
|
||||
offset += len(batch_ids)
|
||||
if len(batch_ids) < page:
|
||||
break
|
||||
return {
|
||||
"ids": ids,
|
||||
"documents": documents,
|
||||
"embeddings": embeddings,
|
||||
"metadatas": metadatas,
|
||||
}
|
||||
|
||||
|
||||
def restore_kb(kb_id: int, snapshot: dict) -> None:
|
||||
"""仅用于重建失败回滚；恢复原 chunk id 和元数据，消费记录仍可对应。"""
|
||||
ids = _as_list(snapshot.get("ids"))
|
||||
if backend.is_disabled():
|
||||
if ids:
|
||||
raise VectorIndexDisabled("本地向量索引已禁用，无法恢复向量快照")
|
||||
return
|
||||
documents = _as_list(snapshot.get("documents"))
|
||||
embeddings = _as_list(snapshot.get("embeddings"))
|
||||
metadatas = _as_list(snapshot.get("metadatas"))
|
||||
if any(len(values) != len(ids) for values in (documents, embeddings, metadatas)) or len(set(ids)) != len(ids):
|
||||
raise ValueError("知识索引快照不完整，已取消恢复")
|
||||
remote = backend.is_qdrant()
|
||||
for identifier, text, vector, meta in zip(ids, documents, embeddings, metadatas):
|
||||
if not isinstance(identifier, str) or not identifier or not isinstance(text, str) or not text.strip() or not isinstance(meta, dict):
|
||||
raise ValueError("知识索引快照正文或来源无效，已取消恢复")
|
||||
if not len(vector) or (remote and len(vector) != 512) or any(not math.isfinite(float(v)) for v in vector):
|
||||
raise ValueError("知识索引快照向量无效，已取消恢复")
|
||||
reset_kb(kb_id)
|
||||
if not ids:
|
||||
return
|
||||
items = [{"id": ids[index], "text": documents[index],
|
||||
"embedding": embeddings[index], "meta": metadatas[index]}
|
||||
for index in range(len(ids))]
|
||||
if not remote:
|
||||
_write_fallback(kb_id, items)
|
||||
ensure_collection(kb_id).add(
|
||||
ids=ids,
|
||||
documents=documents,
|
||||
embeddings=embeddings,
|
||||
metadatas=metadatas,
|
||||
)
|
||||
_invalidate_count_cache(kb_id)
|
||||
@@ -0,0 +1,237 @@
|
||||
"""Small isolated RAG process pool with bounded admission and wall-clock deadlines."""
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
from pathlib import Path
|
||||
import queue
|
||||
import selectors
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
|
||||
|
||||
class RagWorkerError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class RagWorkerBusy(RagWorkerError):
|
||||
pass
|
||||
|
||||
|
||||
class RagWorkerTimeout(RagWorkerError):
|
||||
pass
|
||||
|
||||
|
||||
def _number(name: str, default: float, minimum: float, maximum: float) -> float:
|
||||
try:
|
||||
value = float(os.environ.get(name, default))
|
||||
return min(maximum, max(minimum, value)) if math.isfinite(value) else default
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
class _Worker:
|
||||
def __init__(self, command: list[str], stderr_path: Path | None):
|
||||
self.command = command
|
||||
self.stderr_path = stderr_path
|
||||
self.process: subprocess.Popen | None = None
|
||||
self.last_start_ms = 0.0
|
||||
self.last_reused = False
|
||||
|
||||
def start(self):
|
||||
if self.process is not None and self.process.poll() is None:
|
||||
self.last_start_ms = 0.0
|
||||
self.last_reused = True
|
||||
return self.process
|
||||
self.close()
|
||||
started = time.perf_counter()
|
||||
stderr = subprocess.DEVNULL
|
||||
if self.stderr_path:
|
||||
self.stderr_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
stderr = self.stderr_path.open('ab')
|
||||
try:
|
||||
self.process = subprocess.Popen(
|
||||
self.command, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
|
||||
stderr=stderr, bufsize=0)
|
||||
finally:
|
||||
if hasattr(stderr, 'close'):
|
||||
stderr.close()
|
||||
os.set_blocking(self.process.stdin.fileno(), False)
|
||||
os.set_blocking(self.process.stdout.fileno(), False)
|
||||
self.last_start_ms = round((time.perf_counter() - started) * 1000, 2)
|
||||
self.last_reused = False
|
||||
return self.process
|
||||
|
||||
def close(self):
|
||||
process, self.process = self.process, None
|
||||
if process is None:
|
||||
return
|
||||
try:
|
||||
if process.poll() is None:
|
||||
process.kill()
|
||||
process.wait(timeout=0.2)
|
||||
except (OSError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
finally:
|
||||
for pipe in (process.stdin, process.stdout):
|
||||
if pipe:
|
||||
pipe.close()
|
||||
|
||||
def request(self, payload: dict, deadline: float) -> dict:
|
||||
if time.monotonic() >= deadline:
|
||||
raise RagWorkerTimeout('知识库检索超时，请缩小资料范围后重试。')
|
||||
process = self.start()
|
||||
encoded = memoryview((json.dumps(payload, ensure_ascii=False) + '\n').encode('utf-8'))
|
||||
buffer = bytearray()
|
||||
with selectors.DefaultSelector() as selector:
|
||||
selector.register(process.stdin, selectors.EVENT_WRITE)
|
||||
selector.register(process.stdout, selectors.EVENT_READ)
|
||||
while True:
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
raise RagWorkerTimeout('知识库检索超时，请缩小资料范围后重试。')
|
||||
events = selector.select(remaining)
|
||||
if not events:
|
||||
raise RagWorkerTimeout('知识库检索超时，请缩小资料范围后重试。')
|
||||
for key, _ in events:
|
||||
if key.fileobj is process.stdin:
|
||||
try:
|
||||
sent = os.write(key.fd, encoded[:65536])
|
||||
except BlockingIOError:
|
||||
continue
|
||||
encoded = encoded[sent:]
|
||||
if not encoded:
|
||||
selector.unregister(process.stdin)
|
||||
else:
|
||||
try:
|
||||
part = os.read(key.fd, 65536)
|
||||
except BlockingIOError:
|
||||
continue
|
||||
if not part:
|
||||
raise RagWorkerError('知识库检索进程意外退出，请重试。')
|
||||
buffer.extend(part)
|
||||
if len(buffer) > 32 * 1024 * 1024:
|
||||
raise RagWorkerError('知识库检索返回内容过大，请减少检索片段数。')
|
||||
if b'\n' in buffer:
|
||||
line, trailing = buffer.split(b'\n', 1)
|
||||
if trailing.strip():
|
||||
raise RagWorkerError('知识库检索进程返回格式异常，请重试。')
|
||||
try:
|
||||
response = json.loads(line)
|
||||
except (ValueError, UnicodeDecodeError) as exc:
|
||||
raise RagWorkerError('知识库检索进程返回格式异常，请重试。') from exc
|
||||
if not isinstance(response, dict):
|
||||
raise RagWorkerError('知识库检索进程返回格式异常，请重试。')
|
||||
return response
|
||||
|
||||
|
||||
class RagWorkerPool:
|
||||
"""At most three active workers, a small waiting room, no in-process fallback."""
|
||||
def __init__(self, *, size: int = 2, total_timeout: float = 45,
|
||||
queue_timeout: float = 3, max_pending: int | None = None,
|
||||
command: list[str] | None = None, stderr_path: Path | None = None):
|
||||
self.size = min(3, max(1, int(size)))
|
||||
self.total_timeout = max(0.01, float(total_timeout))
|
||||
self.queue_timeout = max(0.0, min(float(queue_timeout), self.total_timeout))
|
||||
self._closed = False
|
||||
self._workers = [_Worker(command or [sys.executable, str(Path(__file__).with_name('rag_worker.py'))],
|
||||
stderr_path) for _ in range(self.size)]
|
||||
self._available = queue.Queue(maxsize=self.size)
|
||||
for worker in self._workers:
|
||||
self._available.put_nowait(worker)
|
||||
self._admission = threading.BoundedSemaphore(self.size + (max_pending if max_pending is not None else self.size * 2))
|
||||
|
||||
def request(self, payload: dict) -> dict:
|
||||
started = time.monotonic()
|
||||
deadline = started + self.total_timeout
|
||||
if self._closed:
|
||||
raise RagWorkerError('知识库检索服务正在关闭，请稍后重试。')
|
||||
if not self._admission.acquire(blocking=False):
|
||||
raise RagWorkerBusy('知识库检索任务较多，等待队列已满，请稍后重试。')
|
||||
worker = None
|
||||
try:
|
||||
try:
|
||||
worker = self._available.get(timeout=min(self.queue_timeout, max(0, deadline - time.monotonic())))
|
||||
except queue.Empty as exc:
|
||||
raise RagWorkerBusy('知识库检索排队超时，请稍后重试。') from exc
|
||||
queue_ms = round((time.monotonic() - started) * 1000, 2)
|
||||
try:
|
||||
response = worker.request(payload, deadline)
|
||||
if response.get('ok') is not True:
|
||||
from ..system_log import redact
|
||||
detail = redact(response.get('error') or '未知错误', 400)
|
||||
error = RagWorkerError(f'知识库检索失败：{detail}')
|
||||
setattr(error, 'retrieval_metric', response.get('retrieval_metric') or {"error": detail})
|
||||
raise error
|
||||
if payload.get('op') == 'warmup':
|
||||
return response
|
||||
if not isinstance(response.get('results'), list):
|
||||
raise RagWorkerError('知识库检索返回格式异常，请重试。')
|
||||
except Exception as exc:
|
||||
# A timed-out response may never be consumed by a subsequent request.
|
||||
worker.close()
|
||||
if isinstance(exc, RagWorkerError):
|
||||
raise
|
||||
raise RagWorkerError('知识库检索进程通信失败，请重试。') from exc
|
||||
response['transport_timings'] = {
|
||||
'queue_ms': queue_ms, 'total_ms': round((time.monotonic() - started) * 1000, 2),
|
||||
'worker_ms': round((time.monotonic() - started) * 1000 - queue_ms, 2),
|
||||
'process_start_ms': worker.last_start_ms,
|
||||
'process_reused': worker.last_reused,
|
||||
'worker_pid': int(worker.process.pid) if worker.process else 0,
|
||||
}
|
||||
# Child search timing is reported in the response envelope. The
|
||||
# residual after spawn and search is the parent/child pipe and JSON
|
||||
# transport overhead, useful for spotting IPC regressions.
|
||||
child_ms = float((response.get('worker_timings') or {}).get('search_ms') or 0.0)
|
||||
response['transport_timings']['ipc_ms'] = round(max(
|
||||
0.0,
|
||||
response['transport_timings']['worker_ms']
|
||||
- child_ms
|
||||
- float(response['transport_timings']['process_start_ms'] or 0.0)), 2)
|
||||
return response
|
||||
finally:
|
||||
if worker is not None:
|
||||
self._available.put_nowait(worker)
|
||||
self._admission.release()
|
||||
|
||||
def prewarm(self) -> dict:
|
||||
"""Start one resident worker and warm its model/client in the background."""
|
||||
return self.request({'op': 'warmup'})
|
||||
|
||||
def close(self):
|
||||
self._closed = True
|
||||
for worker in self._workers:
|
||||
worker.close()
|
||||
|
||||
|
||||
_pool: RagWorkerPool | None = None
|
||||
_pool_lock = threading.Lock()
|
||||
|
||||
|
||||
def get_pool() -> RagWorkerPool:
|
||||
global _pool
|
||||
with _pool_lock:
|
||||
if _pool is None:
|
||||
from ..config import DATA_DIR
|
||||
_pool = RagWorkerPool(
|
||||
# The local embedding model is loaded once per worker. Keep one
|
||||
# resident worker by default on the single-host deployment to
|
||||
# avoid duplicating several hundred MB and triggering macOS
|
||||
# memory pressure; operators can raise RAG_WORKER_POOL_SIZE
|
||||
# explicitly when concurrent retrieval is required.
|
||||
size=int(_number('RAG_WORKER_POOL_SIZE', 1, 1, 3)),
|
||||
total_timeout=_number('RAG_SEARCH_TIMEOUT_SECONDS', 45, 1, 180),
|
||||
# Content tasks commonly fan out several segment requests at
|
||||
# once. Give queued requests enough time to use the single
|
||||
# memory-efficient worker instead of surfacing a spurious
|
||||
# timeout while an earlier segment is still being searched.
|
||||
queue_timeout=_number('RAG_QUEUE_TIMEOUT_SECONDS', 10, 0, 30),
|
||||
max_pending=int(_number('RAG_MAX_PENDING', 8, 0, 32)),
|
||||
stderr_path=DATA_DIR / 'logs' / 'rag-worker.log')
|
||||
atexit.register(_pool.close)
|
||||
return _pool
|
||||
Reference in New Issue
Block a user