111 lines
4.0 KiB
Python
111 lines
4.0 KiB
Python
"""知识库导入去重：只比较精确正文，不以语义相似度删除资料。"""
|
|
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)
|