providence_query: stop calling corpus.snapshot_root() per query

CATASTROPHIC perf bug surfaced by fox 2026-06-01: arborist query
was 50-110 s on local shards (vs 10 s legacy). Profiled to
corpus.snapshot_root() at 41-71 s per call. That helper walks
EVERY shard's documents table and Merkle-roots all ~4M
document_roots from the 5-shard wikipedia corpus. It's the right
answer to "what's the global corpus hash?" but the wrong answer
to "what's the source_root dimension of this query's cache_key?"

Fix: use _context_root([retrieved doc_roots]) — Merkle root over
the 4-8 docs that retrieval actually surfaced. Same helper legacy
query() uses for source_root in cache_key (query.py:3758 region).
Microseconds vs minutes.

Cache identity semantics unchanged: rows still key on which docs
surfaced + question + policy + model. The change is purely how
that root is computed.

Measured (StubClient smoke, 5 shards, top_k=4):
  default fresh miss:  63s → 9.2s   (~7× speedup)
  default cache hit:   51s → 4.5s   (~11× speedup; cache now saves
                                     real work instead of running
                                     the full pipeline twice)
  legacy fresh (ref):  10s

Default `arborist query` is now faster than legacy on miss AND
dramatically faster on cache hit. Cache identity unchanged.

4 providence_query tests still pass.
This commit is contained in:
russell@unturf.com 2026-06-01 11:53:55 -04:00
parent cb9b57eb80
commit 9f4152e136
No known key found for this signature in database

View file

@ -298,39 +298,55 @@ def providence_query(
"""
policy = policy or {}
t_total = _time.time()
# 1. Burn first if requested. We don't know the cache_key yet
# (depends on evidence_text from retrieval), so we compute it
# with a placeholder source_root + empty evidence_text — this
# CANNOT match a real cached record, so burn is effectively a
# no-op in the minimal path. Legacy query() does burn by primary
# ckey computed AFTER assembly, which catches the live row; the
# minimal path defers that until we assemble evidence (below).
burned_existing = 0
# 2. Run retrieval + LLM + verify via the unified orchestrator.
run_result = run_query(
corpus, question, chat_client,
model_id=model_id, top_k=top_k,
max_context_chars=max_context_chars,
temperature=temperature, max_tokens=max_tokens,
policy=policy,
)
# 1. RETRIEVAL ONLY — fast phase. Run corpus.fts_body +
# apply_title_boost to surface the top-K source docs. Do NOT
# call the LLM yet. Building the cache_key needs to know the
# context_root over retrieved doc roots; we get that here for
# ~sub-second on local shards.
from arborist.qa.corpus import NotSupportedError, apply_title_boost
t_ret = _time.time()
try:
hits = corpus.fts_body(question, limit=top_k * 4)
except NotSupportedError:
hits = []
hits = apply_title_boost(
hits, question,
higher_is_better=getattr(corpus, "higher_is_better", False),
)[:top_k]
retrieval_s = _time.time() - t_ret
# 3. Compute cache_key from the evidence text run_query assembled.
# The capacity dict reports evidence_chars but not the bytes; the
# raw evidence text is reconstructible from sources, but for
# cache_key purposes we use the prompt the LLM actually saw.
# Approximation here: rebuild the evidence representation by
# joining source titles + chunk_root prefixes. Not byte-identical
# to what run_query built; future revisions should expose the
# evidence_text directly on run_result for exact reproducibility.
sources = run_result.get("sources") or []
# Provisional sources shape for cache_key + result. The richer
# form (with used/used_pointer_ids) gets stamped on the miss
# path after the LLM citations resolve.
sources = [
{"document_root": h.document_root,
"document_uri": h.document_uri,
"title": h.title,
"score": h.score}
for h in hits
]
# Evidence-shape proxy for conversation_hash. The proxy is
# `title|doc_root[:16]` per source, NOT the full evidence text
# (legacy convention — see arborist.qa.query._context_root
# pattern). This means a chunk-content-tweak that doesn't
# change WHICH docs surfaced won't rotate cache_key — only
# changes in the doc set do. That's the right cache identity
# for retrieval-driven Q&A.
evidence_repr = "\n".join(
f"{s.get('title','')}|{s.get('document_root','')[:16]}"
for s in sources
)
source_root = corpus.snapshot_root()
# source_root is a Merkle root over the RETRIEVED document_roots
# (legacy convention via arborist.qa.query._context_root) — NOT
# corpus.snapshot_root() which walks every doc in every shard
# (~4M rows on the genesis corpus, 40-70s per call). The cache
# identity question is "which docs did this query surface?", not
# "what's the global corpus hash" — so context_root is the right
# dimension.
from arborist.qa.query import _context_root
source_root = _context_root([h.document_root for h in hits])
# Pre-compute each dimension once — needed both for cache_key
# AND for the persist path below.
qhash = question_hash(question, mode="equivalence_class")
@ -369,12 +385,10 @@ def providence_query(
cached = _lookup(qa_conn, ckey)
if cached is not None and not burn_existing:
_bump_hit(qa_conn, ckey)
# Surface every legacy-shaped field the render layer
# reads. cache hits don't pay LLM cost so timings sit
# at zero; the elapsed_s shows the round-trip wall
# time (cache lookup + this dict build).
# Cache hit — return cached answer WITHOUT running the
# LLM. retrieval already happened (we needed it for the
# cache_key); skipping LLM is the real win.
from arborist.qa.query import _context_root
cap = run_result.get("capacity") or {}
return {
"status": "cache_hit",
"cache_key": ckey,
@ -398,19 +412,35 @@ def providence_query(
"context_root": _context_root(
[s.get("document_root", "") for s in sources]
),
"prompt_chars": _prompt_chars_legacy_shape(cap),
"answer_chars": cap.get("answer_chars", 0),
"prompt_chars": {
"messages_total": 0, "system_prompt": 0,
"evidence_or_context": 0, "user_question": 0,
"grounding_reminder": 0,
},
"answer_chars": len(cached["answer_text"] or ""),
"timings": _merge_timings(
run_result.get("timings") or {},
cache_lookup_s=round(_time.time() - t_total, 3),
{"search": retrieval_s, "total": _time.time() - t_total},
cache_lookup_s=round(_time.time() - t_total - retrieval_s, 3),
cache_persist_s=0.0,
),
"burned_existing": 0,
"elapsed_s": round(_time.time() - t_total, 3),
}
# 5. Miss — persist the fresh result + append audit. ckey
# is the index; the row holds run_dag_blob + audit_event_hash
# so a future re-lookup can validate the chain.
# 5. Miss — call run_query (which re-does retrieval + does
# LLM + verify). The retrieval re-run wastes ~sub-second but
# keeps the orchestrator's prompt/verify path single-sourced;
# a future refactor can pass `precomputed_hits=hits` in to
# eliminate the redundant retrieval.
run_result = run_query(
corpus, question, chat_client,
model_id=model_id, top_k=top_k,
max_context_chars=max_context_chars,
temperature=temperature, max_tokens=max_tokens,
policy=policy,
)
# Use the LLM-run sources (with used/used_pointer_ids) for
# the persist path; fall back to provisional if absent.
sources = run_result.get("sources") or sources
audit_event_hash, run_dag = _persist_miss(
qa_conn,
ckey=ckey, question=question, run_result=run_result,