feat(server): add workbench API, DSH planner worker and RAG pipeline

2026-09-16 19:06:20 +08:00
parent 72fc5b47c7
commit d4b5b3f43a
122 changed files with 31906 additions and 0 deletions
@@ -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)