258 lines
10 KiB
Python
258 lines
10 KiB
Python
"""Persistent, topic-neutral reading records scoped to an organization and library.
|
|
|
|
Only exact current input bytes are reused. File manifests are an optional, conservative
|
|
estimate: missing index revision metadata always means zero estimated reuse.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import time
|
|
import uuid
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
|
|
from sqlmodel import select
|
|
|
|
from .config import DATA_DIR
|
|
from .models import KnowledgeBase, KnowledgeFile
|
|
|
|
KNOWLEDGE_DIR = DATA_DIR / "training-knowledge"
|
|
GENERIC_TARGET = 3600
|
|
CACHE_POLICY = "training-knowledge-ordered-v2"
|
|
LOCK_WAIT_SECONDS = 600
|
|
|
|
|
|
class ReadingCancelled(RuntimeError):
|
|
pass
|
|
|
|
|
|
def check_stop(should_stop):
|
|
if should_stop and should_stop():
|
|
raise ReadingCancelled("资料阅读已取消")
|
|
|
|
|
|
def canonical_entries(entries):
|
|
"""Exclude snapshot E IDs, library display names, topic and selection metadata."""
|
|
return [{"kb_id": item["kb_id"], "file": item["file"], "chunk_id": item["chunk_id"],
|
|
"text": item["text"], "kind": "raw", "source_count": 1,
|
|
"source_seq": item.get("source_seq"), "heading": item.get("heading", "")} for item in entries]
|
|
|
|
|
|
def _directory(organization_id, kb_id):
|
|
for value in (organization_id, kb_id):
|
|
if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
|
|
raise ValueError("知识阅读记录的组织或知识库编号无效")
|
|
root = KNOWLEDGE_DIR.resolve()
|
|
directory = KNOWLEDGE_DIR / str(organization_id) / str(kb_id)
|
|
if directory.resolve() != root / str(organization_id) / str(kb_id):
|
|
raise ValueError("知识阅读记录路径无效")
|
|
return directory
|
|
|
|
|
|
def fingerprint(entries, max_chars, signature):
|
|
payload = {"policy": CACHE_POLICY, "signature": signature, "max_chars": max_chars,
|
|
"entries": canonical_entries(entries)}
|
|
return hashlib.sha256(json.dumps(payload, ensure_ascii=False, sort_keys=True,
|
|
separators=(",", ":")).encode("utf-8")).hexdigest()
|
|
|
|
|
|
def _cache_path(directory, key):
|
|
return directory / key[:2] / f"{key}.json"
|
|
|
|
|
|
def _saved(path, key, max_chars, min_quality="any"):
|
|
try:
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
summary = data.get("text")
|
|
quality = str(data.get("quality") or "full")
|
|
quality_ok = min_quality != "full" or quality == "full"
|
|
if (quality_ok and data.get("fingerprint") == key and isinstance(summary, str)
|
|
and summary.strip() and len(summary) <= max_chars):
|
|
return summary
|
|
except (OSError, ValueError, TypeError, AttributeError):
|
|
pass
|
|
return None
|
|
|
|
|
|
@contextmanager
|
|
def _singleflight(path, should_stop):
|
|
# POSIX flock serializes separate worker processes. Windows keeps the
|
|
# in-process call path working via threading; local production is Linux.
|
|
try:
|
|
import fcntl
|
|
except ImportError:
|
|
fcntl = None
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with path.open("a") as handle:
|
|
try:
|
|
os.chmod(path, 0o600)
|
|
except OSError:
|
|
pass
|
|
deadline = time.monotonic() + LOCK_WAIT_SECONDS
|
|
while True:
|
|
check_stop(should_stop)
|
|
try:
|
|
if fcntl is not None:
|
|
fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
break
|
|
except BlockingIOError:
|
|
if time.monotonic() >= deadline:
|
|
raise TimeoutError("等待同批资料阅读超时，请稍后重试")
|
|
time.sleep(0.1)
|
|
try:
|
|
yield
|
|
finally:
|
|
if fcntl is not None:
|
|
fcntl.flock(handle, fcntl.LOCK_UN)
|
|
|
|
|
|
def reusable_summary(organization_id, kb_id, entries, max_chars, signature, generate,
|
|
should_stop=None, *, quality="full", min_quality="any"):
|
|
"""Return (validated text, reused, fingerprint). Failed/cancelled calls never save."""
|
|
check_stop(should_stop)
|
|
key = fingerprint(entries, max_chars, signature)
|
|
if not signature:
|
|
summary = generate()
|
|
check_stop(should_stop)
|
|
_validate(summary, max_chars)
|
|
return summary.strip(), False, ""
|
|
directory = _directory(organization_id, kb_id)
|
|
path = _cache_path(directory, key)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with _singleflight(path.with_suffix(".lock"), should_stop):
|
|
check_stop(should_stop)
|
|
summary = _saved(path, key, max_chars, min_quality=min_quality)
|
|
if summary is not None:
|
|
return summary, True, key
|
|
summary = generate()
|
|
check_stop(should_stop)
|
|
_validate(summary, max_chars)
|
|
summary = summary.strip()
|
|
staging = path.with_name(f"{key}.{uuid.uuid4().hex}.tmp")
|
|
try:
|
|
with staging.open("x", encoding="utf-8") as handle:
|
|
os.chmod(staging, 0o600)
|
|
json.dump({"fingerprint": key, "text": summary, "quality": quality}, handle, ensure_ascii=False)
|
|
check_stop(should_stop)
|
|
os.replace(staging, path)
|
|
finally:
|
|
staging.unlink(missing_ok=True)
|
|
return summary, False, key
|
|
|
|
|
|
def _validate(summary, max_chars):
|
|
if not isinstance(summary, str) or not summary.strip():
|
|
raise ValueError("分批阅读未返回摘要")
|
|
if len(summary) > max_chars:
|
|
raise ValueError(f"分批摘要超出 {max_chars} 字符预算，已停止避免截断资料")
|
|
|
|
|
|
def index_revision(kb_id):
|
|
"""Stat only; never open/create Chroma or scan indexed text for an estimate."""
|
|
from .rag import store
|
|
from .rag import backend
|
|
if backend.is_qdrant():
|
|
return "qdrant:" + store.kb_revision(kb_id)
|
|
path = store._fallback_path(kb_id)
|
|
try:
|
|
stat = path.stat()
|
|
return f"fallback:{stat.st_dev}:{stat.st_ino}:{stat.st_size}:{stat.st_mtime_ns}:{stat.st_ctime_ns}"
|
|
except OSError:
|
|
# Shared DB/WAL metadata is a conservative global token: any index write
|
|
# invalidates estimates for all libraries. Actual reuse still hashes text.
|
|
from .config import CHROMA_DIR
|
|
paths = [CHROMA_DIR / 'chroma.sqlite3', CHROMA_DIR / 'chroma.sqlite3-wal']
|
|
versions = []
|
|
for i, candidate in enumerate(paths):
|
|
try:
|
|
stat = candidate.stat()
|
|
versions.append([stat.st_dev, stat.st_ino, stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns])
|
|
except FileNotFoundError:
|
|
if i == 0:
|
|
return ''
|
|
versions.append(None)
|
|
except OSError:
|
|
return ''
|
|
return 'chroma:' + hashlib.sha256(json.dumps(versions).encode()).hexdigest()
|
|
|
|
|
|
def record_manifest(organization_id, kb_id, file_name, content_hash, revision, chunk_count,
|
|
signature, index_token, batch_keys):
|
|
signature = f"{CACHE_POLICY}:{signature}" if signature else ""
|
|
if not content_hash or not signature or not index_token or not chunk_count or not batch_keys or not all(batch_keys):
|
|
return
|
|
directory = _directory(organization_id, kb_id)
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
path = directory / "manifest.sqlite"
|
|
with sqlite3.connect(path, timeout=30) as cache:
|
|
os.chmod(path, 0o600)
|
|
cache.execute("""CREATE TABLE IF NOT EXISTS files (
|
|
file_name TEXT NOT NULL, signature TEXT NOT NULL, content_hash TEXT NOT NULL,
|
|
revision INTEGER NOT NULL, chunk_count INTEGER NOT NULL, index_token TEXT NOT NULL,
|
|
batch_keys TEXT NOT NULL, PRIMARY KEY(file_name, signature))""")
|
|
cache.execute("INSERT OR REPLACE INTO files VALUES(?,?,?,?,?,?,?)",
|
|
(file_name, signature, content_hash, revision, chunk_count, index_token, json.dumps(batch_keys)))
|
|
|
|
|
|
def reuse_inventory(session, kb_ids, mapping, signature):
|
|
"""Read-only estimate; requires the same source version and unchanged index file.
|
|
|
|
Caller verifies authorization before entering this helper. Scope is checked again
|
|
and only each current library's own organization directory is considered. A KB
|
|
index rewrite conservatively reports zero until a successful scan refreshes its
|
|
manifests, even when unchanged files will in fact reuse their exact summaries.
|
|
"""
|
|
result = {"reusable_chunks": 0, "reusable_files": 0, "reusable_batches": 0}
|
|
if not signature:
|
|
return result
|
|
signature = f"{CACHE_POLICY}:{signature}"
|
|
from .training_scope import normalize_scope, file_scope
|
|
normalized = normalize_scope(session, kb_ids, mapping or {})
|
|
for kid in dict.fromkeys(kb_ids):
|
|
kb = session.get(KnowledgeBase, kid)
|
|
if not kb or kb.index_status != "ready" or not kb.organization_id:
|
|
continue
|
|
token = index_revision(kid)
|
|
if not token:
|
|
continue
|
|
allowed = file_scope(session, kid, normalized)
|
|
directory = _directory(kb.organization_id, kid)
|
|
path = directory / "manifest.sqlite"
|
|
if not path.is_file():
|
|
continue
|
|
try:
|
|
with sqlite3.connect(f"file:{path}?mode=ro", uri=True) as cache:
|
|
manifests = {row[0]: row[1:] for row in cache.execute(
|
|
"SELECT file_name,content_hash,revision,chunk_count,index_token,batch_keys FROM files WHERE signature=?",
|
|
(signature,))}
|
|
except (OSError, sqlite3.Error):
|
|
continue
|
|
files = session.exec(select(KnowledgeFile).where(KnowledgeFile.kb_id == kid)).all()
|
|
seen = set()
|
|
for file in files:
|
|
if file.file_name in seen or (allowed is not None and file.file_name not in allowed) or not file.content_hash:
|
|
continue
|
|
seen.add(file.file_name)
|
|
saved = manifests.get(file.file_name)
|
|
if not saved or saved[0] != file.content_hash or saved[1] != file.revision or saved[3] != token:
|
|
continue
|
|
try:
|
|
keys = json.loads(saved[4])
|
|
if not isinstance(keys, list) or not keys or not all(
|
|
isinstance(key, str) and len(key) == 64 and all(c in "0123456789abcdef" for c in key)
|
|
and _cache_path(directory, key).is_file() for key in keys
|
|
):
|
|
continue
|
|
count = int(saved[2])
|
|
if count <= 0:
|
|
continue
|
|
except (ValueError, TypeError, OSError):
|
|
continue
|
|
result["reusable_chunks"] += count
|
|
result["reusable_files"] += 1
|
|
result["reusable_batches"] += len(keys)
|
|
return result
|