407 lines
17 KiB
Python
407 lines
17 KiB
Python
"""CLI: python -m tools.knowledge.build (same private .env as the backend)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
from concurrent.futures import ThreadPoolExecutor, wait, FIRST_COMPLETED
|
|
|
|
from tools.knowledge.extraction import extract_batch, extraction_key, ANNOTATION_FIELDS, MAX_BATCH_UNITS, MAX_BATCH_CHARS, ExtractionRejected
|
|
from tools.knowledge.evidence import EXTRACTION_SCHEMA, EXTRACTION_POLICY, STRUCTURAL_KINDS
|
|
from tools.knowledge.sources import PROFILE, profile_units, identifier
|
|
from tools.rag.common import ZPWIKI_ROOT, DATA_DIR, load_zpwiki_page
|
|
from tools.sync.index_state import read_version, durable_json, publication_lock
|
|
|
|
|
|
def batches(units):
|
|
batch, size = [], 0
|
|
for unit in units:
|
|
if batch and (len(batch) >= MAX_BATCH_UNITS or size + len(unit["text"]) > MAX_BATCH_CHARS
|
|
or (unit["source_path"], unit.get("work_label")) != (batch[0]["source_path"], batch[0].get("work_label"))):
|
|
yield batch
|
|
batch, size = [], 0
|
|
batch.append(unit)
|
|
size += len(unit["text"])
|
|
if batch:
|
|
yield batch
|
|
|
|
|
|
def build(
|
|
wiki_root: Path,
|
|
data_dir: Path,
|
|
complete,
|
|
model_identity: dict,
|
|
*,
|
|
graph=True,
|
|
vectors=True,
|
|
workers=2,
|
|
progress=None,
|
|
) -> dict:
|
|
from tools.rag.rebuild_index import acquire_reindex_lock
|
|
|
|
with acquire_reindex_lock(data_dir / "knowledge/.build.lock"):
|
|
return _build(
|
|
wiki_root,
|
|
data_dir,
|
|
complete,
|
|
model_identity,
|
|
graph=graph,
|
|
vectors=vectors,
|
|
workers=workers,
|
|
progress=progress,
|
|
)
|
|
|
|
|
|
def _build(wiki_root, data_dir, complete, model_identity, *, graph, vectors, workers, progress):
|
|
import tempfile
|
|
import shutil
|
|
from tools.sync.index_state import index_snapshot
|
|
from tools.knowledge.publication import checkpoint_state, publish_checkpoint, PENDING_FILE
|
|
from tools.knowledge.store import create_constraints
|
|
from tools.graphrag.graph_db import create_neo4j_driver, get_neo4j_settings
|
|
from tools.sync.sync_indexes import source_files
|
|
with index_snapshot(data_dir) as version:
|
|
if version is None:
|
|
raise RuntimeError("Build the shared source indexes before content extraction.")
|
|
index, summary = prepare_index(wiki_root, data_dir, complete, model_identity,
|
|
graph=graph, vectors=vectors, workers=workers, progress=progress, source_version=version)
|
|
stage = Path(tempfile.mkdtemp(prefix=".knowledge-stage-", dir=data_dir))
|
|
try:
|
|
checkpoint_state(stage, index)
|
|
with publication_lock(data_dir, exclusive=True):
|
|
source_digest=hashlib.sha256(json.dumps(source_files(wiki_root),sort_keys=True).encode()).hexdigest()
|
|
if read_version(data_dir) != version or source_digest != version['source_sha256']:
|
|
raise RuntimeError("Source changed during extraction; checkpoint not activated.")
|
|
if graph:
|
|
settings = get_neo4j_settings()
|
|
with create_neo4j_driver(settings) as driver, driver.session(database=settings.database) as session:
|
|
driver.verify_connectivity()
|
|
create_constraints(session)
|
|
publish_checkpoint(data_dir, stage, index, session)
|
|
else:
|
|
publish_checkpoint(data_dir, stage, index)
|
|
return summary
|
|
finally:
|
|
if not (data_dir / PENDING_FILE).exists():
|
|
shutil.rmtree(stage)
|
|
|
|
|
|
def prepare_index(wiki_root, data_dir, complete, model_identity, *, graph=True, vectors=True, workers=2, progress=None, source_version=None):
|
|
from tools.sync.sync_indexes import source_files
|
|
|
|
version = source_version or read_version(data_dir)
|
|
if not version:
|
|
raise RuntimeError("Build the shared source indexes before content extraction.")
|
|
sources = source_files(wiki_root)
|
|
digest = hashlib.sha256(json.dumps(sources, sort_keys=True).encode()).hexdigest()
|
|
if digest != version["source_sha256"]:
|
|
raise RuntimeError("Wiki differs from the active index; synchronize first.")
|
|
documents = []
|
|
for path in sources:
|
|
if PROFILE.fullmatch(path):
|
|
document = load_zpwiki_page(wiki_root / path, wiki_root=wiki_root)
|
|
raw = (wiki_root / path).read_text()
|
|
document["body_line_offset"] = (
|
|
raw[: raw.index(document["content"])].count("\n") if document["content"] else 0
|
|
)
|
|
documents.append(document)
|
|
documents = [d for d in documents if d.get("published") is not False]
|
|
if not documents:
|
|
raise RuntimeError("No published student profiles to extract.")
|
|
units = [unit for d in documents for unit in profile_units(d, include_records=True)]
|
|
# Keep whole statements, including trailing negations/qualifications.
|
|
# Over-budget source units are retained with an explicit uncertain decision.
|
|
previous = data_dir / "knowledge/index.json"
|
|
prior = json.loads(previous.read_text()).get("units", []) if previous.exists() else []
|
|
|
|
def source_identity(unit):
|
|
return (unit["source_path"], unit.get("work_label"), unit["line"], unit["text"])
|
|
|
|
prior_ids = {source_identity(u): u["id"] for u in prior}
|
|
# Inserting a status note must not invalidate every subsequent profile's
|
|
# model cache just because its position in the corpus has shifted.
|
|
bounded = [
|
|
{
|
|
**unit,
|
|
"id": prior_ids.get(source_identity(unit), identifier(unit["id"], i, unit["text"])),
|
|
}
|
|
for i, unit in enumerate(units)
|
|
]
|
|
local_kinds = STRUCTURAL_KINDS | {"lifecycle"}
|
|
annotatable = [u for u in bounded if u["kind"] not in local_kinds]
|
|
cache = data_dir / "knowledge/cache"
|
|
cache.mkdir(parents=True, exist_ok=True)
|
|
# Every cached decision is bound to this policy, verifier, complete source
|
|
# scope, actor, dates and model identity. Legacy omissions are never reused.
|
|
proposals = {}
|
|
for file in cache.glob("*.json"):
|
|
try:
|
|
previous = json.loads(file.read_text())
|
|
if (not isinstance(previous,list) or not previous
|
|
or any(row.get("annotation_policy") != EXTRACTION_POLICY
|
|
or row.get("annotation_status") not in {"knowledge","irrelevant","uncertain"}
|
|
for row in previous)
|
|
or file.stem != extraction_key(previous,model_identity)):
|
|
continue
|
|
except (ValueError, OSError, AttributeError, TypeError):
|
|
continue
|
|
for row in previous:
|
|
proposals[extraction_key([row],model_identity)] = {k:v for k,v in row.items() if k in ANNOTATION_FIELDS}
|
|
import threading
|
|
requests = 0
|
|
request_lock = threading.Lock()
|
|
def counted_complete(messages):
|
|
nonlocal requests
|
|
with request_lock:
|
|
requests += 1
|
|
return complete(messages)
|
|
|
|
def extract_complete_batch(batch):
|
|
try:
|
|
return extract_batch(batch,counted_complete)
|
|
except (ExtractionRejected, RuntimeError) as error:
|
|
if isinstance(error, RuntimeError) and str(error) != "Incomplete answer model output":
|
|
raise
|
|
if len(batch) == 1:
|
|
if isinstance(error, ExtractionRejected):
|
|
raise ExtractionRejected([f"{batch[0]['source_path']}:{batch[0]['line']}",str(error)]) from error
|
|
raise
|
|
# Every child still needs a decision for every input. A missing ID
|
|
# can never turn into an empty successful annotation.
|
|
middle = len(batch)//2
|
|
return [*extract_complete_batch(batch[:middle]),*extract_complete_batch(batch[middle:])]
|
|
|
|
def process(batch):
|
|
key = extraction_key(batch, model_identity)
|
|
file = cache / (key + ".json")
|
|
if file.exists():
|
|
rows = json.loads(file.read_text())
|
|
# Cache is private and keyed by complete source, prompt and model.
|
|
if ([r["id"] for r in rows] != [r["id"] for r in batch]
|
|
or any(r.get("annotation_policy") != EXTRACTION_POLICY or r.get("annotation_status") not in {"knowledge","irrelevant","uncertain"}
|
|
or extraction_key([r],model_identity) != extraction_key([u],model_identity)
|
|
for r,u in zip(rows,batch))):
|
|
raise RuntimeError("Invalid extraction cache identity.")
|
|
rows = [
|
|
{**unit, **{k: v for k, v in row.items() if k in ANNOTATION_FIELDS}}
|
|
for unit, row in zip(batch, rows)
|
|
]
|
|
cached = True
|
|
else:
|
|
missing = [u for u in batch if extraction_key([u], model_identity) not in proposals]
|
|
fresh = {u["id"]: u for u in extract_complete_batch(missing)} if missing else {}
|
|
rows = [
|
|
fresh.get(u["id"]) or {**u, **proposals[extraction_key([u], model_identity)]}
|
|
for u in batch
|
|
]
|
|
durable_json(file, rows)
|
|
cached = not missing
|
|
return rows, cached
|
|
|
|
extracted, calls, hits = [], 0, 0
|
|
if not 1 <= workers <= 4:
|
|
raise ValueError("Extraction workers must be between 1 and 4.")
|
|
with ThreadPoolExecutor(max_workers=workers) as pool:
|
|
source_batches = iter(batches(annotatable))
|
|
pending = set()
|
|
for _ in range(workers):
|
|
if (batch := next(source_batches, None)) is not None:
|
|
pending.add(pool.submit(process, batch))
|
|
while pending:
|
|
done, pending = wait(pending, return_when=FIRST_COMPLETED)
|
|
# At most `workers` requests are in flight. A slow batch does not
|
|
# block starting the next one; failures never queue the whole corpus.
|
|
completed = [future.result() for future in done]
|
|
for rows, cached in completed:
|
|
extracted.extend(rows)
|
|
hits += int(cached)
|
|
calls += int(not cached)
|
|
if progress:
|
|
progress(calls + hits)
|
|
if (batch := next(source_batches, None)) is not None:
|
|
pending.add(pool.submit(process, batch))
|
|
extracted.extend(
|
|
{**u, "concepts": [], "relations": [], "annotation_status": "structural",
|
|
"annotation_reason": "Literal structural, role, reference or lifecycle record",
|
|
"annotation_policy": EXTRACTION_POLICY}
|
|
for u in bounded if u["kind"] in local_kinds
|
|
)
|
|
order = {u["id"]: i for i, u in enumerate(bounded)}
|
|
extracted.sort(key=lambda u: order[u["id"]])
|
|
from tools.knowledge.academic import completion_event
|
|
|
|
for unit in extracted:
|
|
if event := completion_event(unit):
|
|
unit.update(
|
|
completion_status=event["status"],
|
|
completion_scope=event["scope"],
|
|
completion_event_year=event["event_year"],
|
|
completion_event_date=event["event_date"],
|
|
)
|
|
index = {
|
|
"schema": 1,
|
|
"academic_schema": 1,
|
|
"extraction_schema": EXTRACTION_SCHEMA,
|
|
"extraction_policy": EXTRACTION_POLICY,
|
|
"index_version": version,
|
|
"model": model_identity,
|
|
"source_scope": "published_student_profiles",
|
|
"profiles": len(documents),
|
|
"units": extracted,
|
|
"graph_ready": graph,
|
|
}
|
|
from collections import Counter, defaultdict
|
|
decision_counts = Counter(u["annotation_status"] for u in extracted)
|
|
per_scope = defaultdict(Counter)
|
|
for unit in extracted:
|
|
per_scope[(unit["person"],unit.get("work_label"),unit["phase"])][unit["annotation_status"]] += 1
|
|
index["coverage"] = {"decisions":dict(decision_counts), "units":len(extracted),
|
|
"by_person_work_phase":[{"person":key[0],"work_label":key[1],"phase":key[2],"decisions":dict(value)}
|
|
for key,value in sorted(per_scope.items(),key=lambda pair:str(pair[0]))]}
|
|
if len(extracted) != len(bounded) or {u["id"] for u in extracted} != {u["id"] for u in bounded}:
|
|
raise RuntimeError("Missing or duplicate decisions in extraction checkpoint")
|
|
if vectors:
|
|
index.update(embedding_index(extracted, data_dir))
|
|
index["version"] = identifier(json.dumps(index, ensure_ascii=False, sort_keys=True))
|
|
if source_files(wiki_root) != sources:
|
|
raise RuntimeError("Source changed during extraction; checkpoint not activated.")
|
|
return index, {
|
|
|
|
"profiles": len(documents),
|
|
"evidence_units": len(extracted),
|
|
"llm_batches": calls,
|
|
"cached_batches": hits,
|
|
"model_requests": requests,
|
|
"decisions": dict(decision_counts),
|
|
"uncertain_relations": sum(r.get("uncertain_relations",0) for r in extracted),
|
|
"concepts": len({c.casefold() for r in extracted for c in r["concepts"]}),
|
|
"discarded_concepts": sum(r.get("discarded_concepts", 0) for r in extracted),
|
|
"discarded_relations": sum(r.get("discarded_relations", 0) for r in extracted),
|
|
"version": index["version"],
|
|
"graph_ready": graph,
|
|
}
|
|
|
|
|
|
def embedding_index(units: list[dict], data_dir: Path) -> dict:
|
|
import base64
|
|
import zlib
|
|
import numpy as np
|
|
from tools.rag.embedding_cache import EmbeddingCache, cache_key
|
|
from tools.rag.embedding_utils import (
|
|
embedding_model_name,
|
|
model_profile,
|
|
split_passage,
|
|
embed_passages,
|
|
)
|
|
|
|
profile = model_profile(embedding_model_name())
|
|
cache = EmbeddingCache(data_dir / "knowledge/embedding_cache.sqlite")
|
|
inputs = [split_passage(u["text"]) for u in units]
|
|
keys = [[cache_key(profile, s) for s in pieces] for pieces in inputs]
|
|
found = cache.load([key for group in keys for key in group])
|
|
shared = data_dir / "embedding_cache.sqlite"
|
|
if shared.exists():
|
|
import sqlite3
|
|
|
|
try:
|
|
found.update(
|
|
EmbeddingCache(shared, readonly=True).load(
|
|
[key for group in keys for key in group if key not in found]
|
|
)
|
|
)
|
|
except sqlite3.Error:
|
|
pass # A cache is an optimization, never the source of truth.
|
|
missing = {
|
|
key: text
|
|
for group, pieces in zip(keys, inputs)
|
|
for key, text in zip(group, pieces)
|
|
if key not in found
|
|
}
|
|
if missing:
|
|
fresh = dict(zip(missing, embed_passages(list(missing.values()))))
|
|
cache.save(fresh)
|
|
found.update(fresh)
|
|
vectors = []
|
|
for group in keys:
|
|
value = np.mean([found[key] for key in group], axis=0)
|
|
value /= np.linalg.norm(value)
|
|
vectors.append(value)
|
|
matrix = np.asarray(vectors, dtype="<f4")
|
|
return {
|
|
"embedding_fingerprint": profile["fingerprint"],
|
|
"vector_shape": list(matrix.shape),
|
|
"vectors": base64.b64encode(zlib.compress(matrix.tobytes())).decode("ascii"),
|
|
}
|
|
|
|
|
|
def configured_extractor():
|
|
from functools import partial
|
|
from tools.answers.model_client import complete_with_configured_model
|
|
|
|
provider = os.getenv("ANSWER_PROVIDER", "openai_compatible")
|
|
legacy = bool(os.getenv("OPENWEBUI_API_KEY"))
|
|
identity = {
|
|
"provider": provider,
|
|
"model": os.getenv("ANSWER_MODEL")
|
|
or ("model120-fast" if provider == "openai_compatible" and legacy else ""),
|
|
"endpoint": os.getenv("ANSWER_LLM_URL")
|
|
or (
|
|
"https://api.anthropic.com/v1/messages"
|
|
if provider == "anthropic"
|
|
else "https://ui.tukekemt.xyz/api/chat/completions" if legacy else ""
|
|
),
|
|
"max_tokens": 4096,
|
|
}
|
|
return partial(complete_with_configured_model, max_tokens=4096), identity
|
|
|
|
|
|
def refresh_existing(wiki_root: Path, data_dir: Path) -> dict | None:
|
|
"""Once enabled by an initial build, keep content aligned on later syncs.
|
|
|
|
Cached unchanged inputs need no LLM calls. A failed refresh leaves content
|
|
unavailable instead of serving the previous snapshot as current evidence.
|
|
"""
|
|
if not (data_dir / "knowledge/index.json").exists():
|
|
return None
|
|
complete, identity = configured_extractor()
|
|
return build(wiki_root, data_dir, complete, identity)
|
|
|
|
|
|
def main():
|
|
# Dotenv is intentionally local to the CLI; services already load env.
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv(Path(__file__).resolve().parents[2] / ".env", override=False)
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
"--wiki", type=Path, default=Path(os.getenv("ZPWIKI_ROOT", str(ZPWIKI_ROOT)))
|
|
)
|
|
parser.add_argument("--data", type=Path, default=DATA_DIR)
|
|
parser.add_argument("--workers", type=int, choices=range(1,5), default=2)
|
|
args = parser.parse_args()
|
|
complete, identity = configured_extractor()
|
|
print(
|
|
json.dumps(
|
|
build(
|
|
args.wiki,
|
|
args.data,
|
|
complete,
|
|
identity,
|
|
workers=args.workers,
|
|
progress=lambda count: (
|
|
print(f"Validated extraction batches: {count}", flush=True)
|
|
if count % 10 == 0
|
|
else None
|
|
),
|
|
),
|
|
ensure_ascii=False,
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|