351 lines
13 KiB
Python
351 lines
13 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
|
|
from itertools import islice
|
|
|
|
from tools.knowledge.extraction import extract_batch, extraction_key, PROMPT
|
|
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) >= 64 or size + len(unit["text"]) > 16000):
|
|
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):
|
|
from tools.sync.sync_indexes import source_files
|
|
|
|
version = 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.
|
|
# Only the model's annotation view may be shortened; evidence never is.
|
|
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 = {"reference", "lifecycle", "work_record", "meeting", "consultant"}
|
|
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)
|
|
# Reuse annotations independently of dates/line offsets. These are parsed
|
|
# locally, not inferred by the LLM. Legacy checkpoints are accepted only
|
|
# when their complete prompt/model/input fingerprint verifies.
|
|
proposal_fields = {
|
|
"concepts",
|
|
"relations",
|
|
"discarded_concepts",
|
|
"discarded_relations",
|
|
"annotation_input_limited",
|
|
}
|
|
proposals = {}
|
|
for file in cache.glob("*.json"):
|
|
previous = json.loads(file.read_text())
|
|
inputs = [{k: v for k, v in row.items() if k not in proposal_fields} for row in previous]
|
|
legacy = identifier(
|
|
PROMPT,
|
|
json.dumps(inputs, sort_keys=True, ensure_ascii=False),
|
|
json.dumps(model_identity, sort_keys=True),
|
|
)
|
|
if file.stem not in {legacy, extraction_key(inputs, model_identity)}:
|
|
continue
|
|
for row in previous:
|
|
proposals[extraction_key([row], model_identity)] = {
|
|
k: v for k, v in row.items() if k in proposal_fields
|
|
}
|
|
|
|
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]:
|
|
raise RuntimeError("Invalid extraction cache identity.")
|
|
rows = [
|
|
{**unit, **{k: v for k, v in row.items() if k in proposal_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_batch(missing, complete)} 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:
|
|
pending = iter(batches(annotatable))
|
|
while window := list(islice(pending, workers)):
|
|
# Bounded submission: an API failure cannot leave hundreds of
|
|
# queued requests running after a failed extraction.
|
|
for rows, cached in pool.map(process, window):
|
|
extracted.extend(rows)
|
|
hits += int(cached)
|
|
calls += int(not cached)
|
|
if progress:
|
|
progress(calls + hits)
|
|
# A batch can straddle two profiles. Give profiles with no accepted
|
|
# annotation their own bounded pass rather than silently losing them.
|
|
absent = {u["person"] for u in extracted} - {u["person"] for u in extracted if u["concepts"]}
|
|
mixed = {
|
|
u["person"]
|
|
for batch in batches(annotatable)
|
|
if len({r["person"] for r in batch}) > 1
|
|
for u in batch
|
|
}
|
|
absent &= mixed
|
|
by_id = {u["id"]: u for u in extracted}
|
|
for person in sorted(absent):
|
|
own = [u for u in annotatable if u["person"] == person and u["kind"] != "supervisor"][:32]
|
|
if own:
|
|
rows, cached = process(own)
|
|
by_id.update({u["id"]: u for u in rows})
|
|
calls += int(not cached)
|
|
hits += int(cached)
|
|
extracted = [by_id[u["id"]] for u in extracted]
|
|
extracted.extend(
|
|
{**u, "concepts": [], "relations": []} 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,
|
|
"index_version": version,
|
|
"model": model_identity,
|
|
"source_scope": "published_student_profiles",
|
|
"profiles": len(documents),
|
|
"units": extracted,
|
|
"graph_ready": graph,
|
|
}
|
|
if vectors:
|
|
index.update(embedding_index(extracted, data_dir))
|
|
index["version"] = identifier(json.dumps(index, ensure_ascii=False, sort_keys=True))
|
|
with publication_lock(data_dir, exclusive=True):
|
|
if read_version(data_dir) != version or source_files(wiki_root) != sources:
|
|
raise RuntimeError("Source changed during extraction; checkpoint not activated.")
|
|
if graph:
|
|
from tools.knowledge.store import write_graph
|
|
|
|
write_graph(index)
|
|
durable_json(data_dir / "knowledge/index.json", index)
|
|
return {
|
|
"profiles": len(documents),
|
|
"evidence_units": len(extracted),
|
|
"llm_batches": calls,
|
|
"cached_batches": hits,
|
|
"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)
|
|
args = parser.parse_args()
|
|
complete, identity = configured_extractor()
|
|
print(
|
|
json.dumps(
|
|
build(
|
|
args.wiki,
|
|
args.data,
|
|
complete,
|
|
identity,
|
|
progress=lambda count: (
|
|
print(f"Validated extraction batches: {count}", flush=True)
|
|
if count % 10 == 0
|
|
else None
|
|
),
|
|
),
|
|
ensure_ascii=False,
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|