feat(server): add workbench API, DSH planner worker and RAG pipeline
This commit is contained in:
yi
@@ -0,0 +1,257 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user