diff --git a/aborist/qa/query.py b/aborist/qa/query.py index ad5e7d1..308d0f8 100644 --- a/aborist/qa/query.py +++ b/aborist/qa/query.py @@ -309,16 +309,25 @@ def query( ) -> dict: """Answer `question` using the corpus. Cache to qa_db. Returns a result dict.""" policy = policy or DEFAULT_QUERY_POLICY + t_start = time.monotonic() # 1. Search. + t_phase = time.monotonic() hits = _search_corpus(shards_dir, single_db, question, over_fetch) if not hits: - return {"status": "no_sources", "msg": "FTS5 search returned no hits"} - # 1b. Rerank — title-boost gives topically-named docs priority over - # body-only matches like Tell_(poker) for "tell me about ...". + return { + "status": "no_sources", + "msg": "FTS5 search returned no hits", + "timings": { + "search_ms": _ms_since(t_phase), + "total_ms": _ms_since(t_start), + }, + } hits = _rerank(hits, question) + search_ms = _ms_since(t_phase) # 2. Pull doc texts within budget. + t_phase = time.monotonic() chosen: list[_Hit] = [] context_parts: list[str] = [] char_budget = max_context_chars @@ -337,9 +346,17 @@ def query( break if not chosen: - return {"status": "no_sources", "msg": "top-k hits had cold or empty content"} - + return { + "status": "no_sources", + "msg": "top-k hits had cold or empty content", + "timings": { + "search_ms": search_ms, + "context_ms": _ms_since(t_phase), + "total_ms": _ms_since(t_start), + }, + } context = "\n\n".join(context_parts) + context_ms = _ms_since(t_phase) # 3. Build messages + hashes. messages = [ @@ -368,11 +385,13 @@ def query( qa_conn = connect(qa_db) try: # 4. Cache lookup. + t_phase = time.monotonic() cached = qa_conn.execute( "SELECT * FROM providence_cache " "WHERE cache_key = ? AND falsification_state = 'live'", (ckey,), ).fetchone() + cache_lookup_ms = _ms_since(t_phase) if cached is not None: now = int(time.time()) with transaction(qa_conn): @@ -389,9 +408,18 @@ def query( "context_root": context_root, "answer_text": cached["answer_text"], "sources": json.loads(cached["merkle_proof"])["sources"], + "timings": { + "search_ms": search_ms, + "context_ms": context_ms, + "cache_lookup_ms": cache_lookup_ms, + "llm_ms": None, + "persist_ms": None, + "total_ms": _ms_since(t_start), + }, } # 5. Cache miss — call LLM. + t_phase = time.monotonic() answer_text = chat_client.chat_completion( messages, model=model_id, @@ -399,8 +427,10 @@ def query( max_tokens=policy["max_tokens"], top_p=policy.get("top_p", 1.0), ) + llm_ms = _ms_since(t_phase) # 6. Persist record + audit event. + t_phase = time.monotonic() proof_obj = { "context_root": context_root, "sources": [ @@ -463,6 +493,7 @@ def query( finally: qa_conn.close() + persist_ms = _ms_since(t_phase) return { "status": "cache_miss_then_written", "audit_mode": "STRICT", @@ -470,4 +501,17 @@ def query( "context_root": context_root, "answer_text": answer_text, "sources": proof_obj["sources"], + "timings": { + "search_ms": search_ms, + "context_ms": context_ms, + "cache_lookup_ms": cache_lookup_ms, + "llm_ms": llm_ms, + "persist_ms": persist_ms, + "total_ms": _ms_since(t_start), + }, } + + +def _ms_since(t: float) -> float: + """Wall-time elapsed in milliseconds, rounded to 1 decimal.""" + return round((time.monotonic() - t) * 1000, 1) diff --git a/aborist/qa/runner.py b/aborist/qa/runner.py index 49a4ef6..e0a1877 100644 --- a/aborist/qa/runner.py +++ b/aborist/qa/runner.py @@ -42,6 +42,10 @@ DEFAULT_POLICY = { } +def _ms_since(t: float) -> float: + return round((time.monotonic() - t) * 1000, 1) + + def ask( conn: sqlite3.Connection, *, @@ -56,6 +60,7 @@ def ask( ) -> dict: """Look up cached answer or run inference. Returns a result dict.""" policy = policy or DEFAULT_POLICY + t_start = time.monotonic() doc = conn.execute( "SELECT document_uri, chunking_version FROM documents " @@ -103,11 +108,13 @@ def ask( doc["chunking_version"], ) + t_lookup = time.monotonic() cached = conn.execute( "SELECT * FROM providence_cache " "WHERE cache_key = ? AND falsification_state = 'live'", (ckey,), ).fetchone() + cache_lookup_ms = _ms_since(t_lookup) if cached is not None: with transaction(conn): now = int(time.time()) @@ -124,8 +131,14 @@ def ask( "source_root": document_root, "answer_text": cached["answer_text"], "merkle_proof": json.loads(cached["merkle_proof"]), + "timings": { + "cache_lookup_ms": cache_lookup_ms, + "llm_ms": None, + "total_ms": _ms_since(t_start), + }, } + t_llm = time.monotonic() answer_text = client.chat_completion( messages, model=model_id, @@ -133,6 +146,7 @@ def ask( max_tokens=policy["max_tokens"], top_p=policy.get("top_p", 1.0), ) + llm_ms = _ms_since(t_llm) leaves = [bytes.fromhex(r["leaf_hash"]) for r in chunk_rows] tree = MerkleTree.build(leaves) @@ -193,4 +207,9 @@ def ask( "source_root": document_root, "answer_text": answer_text, "merkle_proof": proof_obj, + "timings": { + "cache_lookup_ms": cache_lookup_ms, + "llm_ms": llm_ms, + "total_ms": _ms_since(t_start), + }, }