precompute_rag_cache.py

"""Offline warm-up of the pathway report cache.

Retrieval and evidence gating are deterministic in (path, corpus), but a full
run costs minutes of CPU per path. This script computes every pathway report
once, offline, so the running app serves them from cache instead of recomputing
an identical answer on each request.

Run it after the knowledge base changes (the cache key carries a corpus
fingerprint, so stale entries are ignored automatically rather than served).

    python precompute_rag_cache.py            # all reachable targets
    python precompute_rag_cache.py As2O3 PbO  # only these targets

Interrupting is safe: finished paths stay cached, so a re-run resumes.
"""
from __future__ import annotations

import os
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# Some network node names carry zero-width characters, and the tracer prints the
# target name on a miss. On a GBK console that print raises UnicodeEncodeError
# and aborts the run, so force UTF-8 output and never fail on an odd glyph.
for _stream in (sys.stdout, sys.stderr):
    try:
        _stream.reconfigure(encoding="utf-8", errors="replace")
    except (AttributeError, ValueError):
        pass


def main(argv: list[str]) -> int:
    from search_subgraph import CRNTracer, GRAPHML_PATH
    from tracernet.crn import pathways as crn_pathways
    from tracernet.crn.repository import CRNRepository
    from rag_module import RAGService

    project = os.path.dirname(os.path.abspath(__file__))
    repository = CRNRepository.discover([project])
    tracer = CRNTracer(GRAPHML_PATH, crn_repository=repository)
    graph = tracer.G
    service = RAGService(repository)
    print(f"[init] graph={graph.number_of_nodes()} nodes, "
          f"corpus signature={service._report_cache_signature()}", flush=True)

    targets = argv or sorted(graph.nodes())
    print(f"[plan] tracing {len(targets)} candidate targets ...", flush=True)

    # Collect unique paths across every target. The same path is reachable from
    # several targets, and computing it twice would cost minutes for nothing.
    jobs: dict[str, tuple[str, object]] = {}
    for target in targets:
        try:
            _resolved, _graph, _report, all_paths = crn_pathways.trace_material_pathways(
                tracer, graph, target, max_depth=10, limit=100
            )
        except Exception as error:  # a target with no upstream path is normal
            print(f"[trace] {target}: skipped ({type(error).__name__})", flush=True)
            continue
        for info in all_paths or []:
            path_str = info.get("path_str")
            if path_str and path_str not in jobs:
                jobs[path_str] = (info.get("source"), info.get("full_path"))
    print(f"[plan] {len(jobs)} unique pathways to warm", flush=True)

    started = time.time()
    done = 0
    for index, (path_str, (source, full_path)) in enumerate(jobs.items(), 1):
        step = time.time()
        try:
            # Evidence only: the narrative is generated per request, so warming
            # it here would be discarded work -- and with an API key configured
            # it would spend one LLM call per path for nothing.
            service.cached_path_evidence(source, path_str, full_path)
            done += 1
            status = "ok"
        except Exception as error:
            status = f"FAILED {type(error).__name__}: {error}"
        elapsed = time.time() - step
        total = time.time() - started
        print(f"[{index}/{len(jobs)}] {elapsed:6.1f}s  total {total/60:5.1f}m  "
              f"{status}  {path_str[:90]}", flush=True)

    print(f"[done] warmed {done}/{len(jobs)} pathways in "
          f"{(time.time()-started)/60:.1f} minutes", flush=True)
    print(f"[done] cache directory: {service._report_cache_dir()}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))