229 lines
12 KiB
Python
229 lines
12 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)]
|
|
# 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)]
|
|
annotatable=[u for u in bounded if u['kind'] not in {'reference','lifecycle'}]
|
|
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 {'reference','lifecycle'})
|
|
order={u['id']:i for i,u in enumerate(bounded)}
|
|
extracted.sort(key=lambda u:order[u['id']])
|
|
index = {'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()
|