arborist/aborist/qa/runner.py

"""Q&A runner: cache-first lookup -> inference fallback -> provable record.

Implements the v9.8 admissibility invariant:
    No record reused unless all 8 cache_key dimensions match AND state
    is 'live' (not failed/stale/quarantined).

Cache hit  -> STRICT-mode answer with stored Merkle proof.
Cache miss -> call ChatClient, store record, audit event.
"""

from __future__ import annotations

import json
import sqlite3
import time

from aborist import (
    CANONICALIZATION_VERSION,
    SCHEMA_VERSION,
)
from aborist.merkle import MerkleTree, proof_to_dict
from aborist.qa.client import ChatClient
from aborist.qa.keys import (
    cache_key,
    conversation_hash,
    governance_policy_hash,
    model_profile_hash,
    question_hash,
)
from aborist.store import append_audit, transaction


DEFAULT_POLICY = {
    "system_prompt": (
        "Answer the user's question based ONLY on the document below. "
        "If the answer is not in the document, say 'I don't know.' "
        "Do not speculate. Cite the relevant sentence(s)."
    ),
    "temperature": 0.1,
    "top_p": 1.0,
    "max_tokens": 512,
}


def _ms_since(t: float) -> float:
    return round((time.monotonic() - t) * 1000, 1)


def ask(
    conn: sqlite3.Connection,
    *,
    document_root: str,
    question: str,
    client: ChatClient,
    model_id: str,
    revision: str = "",
    quantization: str = "",
    policy: dict | None = None,
    chain: str = "private",
) -> 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 "
        "WHERE document_root = ?",
        (document_root,),
    ).fetchone()
    if doc is None:
        return {"status": "unknown_document"}

    chunk_rows = conn.execute(
        "SELECT idx, leaf_hash, content FROM chunks "
        "WHERE document_root = ? ORDER BY idx ASC",
        (document_root,),
    ).fetchall()
    if not chunk_rows:
        return {"status": "unknown_document"}
    if any(r["content"] is None for r in chunk_rows):
        return {"status": "source_cold", "msg": "rehydrate before asking"}

    document_text = "\n\n".join(r["content"] for r in chunk_rows)

    messages = [
        {"role": "system", "content": policy["system_prompt"]},
        {
            "role": "user",
            "content": (
                f"Document:\n\n{document_text}\n\n---\n\n"
                f"Question: {question}"
            ),
        },
    ]

    qhash = question_hash(question)
    mhash = model_profile_hash(model_id, revision, quantization)
    chash = conversation_hash(messages)
    ghash = governance_policy_hash(policy)
    ckey = cache_key(
        document_root,
        qhash,
        mhash,
        chash,
        ghash,
        SCHEMA_VERSION,
        CANONICALIZATION_VERSION,
        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())
            conn.execute(
                "UPDATE providence_cache "
                "SET hit_count = hit_count + 1, last_hit_at = ? "
                "WHERE cache_key = ?",
                (now, ckey),
            )
        return {
            "status": "cache_hit",
            "audit_mode": "STRICT",
            "cache_key": ckey,
            "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,
        temperature=policy["temperature"],
        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)
    proof_obj = {
        "document_root": document_root,
        "chunk_0_proof": proof_to_dict(tree.proof(0)),
    }
    proof_blob = json.dumps(proof_obj, separators=(",", ":"))

    now = int(time.time())
    with transaction(conn):
        event_hash = append_audit(
            conn,
            event_type="providence_write",
            subject_root=ckey,
            body={
                "source_root": document_root,
                "model_id": model_id,
                "revision": revision,
                "quantization": quantization,
                "chunks_in_context": len(chunk_rows),
                "answer_chars": len(answer_text),
            },
            ts=now,
        )
        conn.execute(
            "INSERT INTO providence_cache "
            "(cache_key, source_root, document_uri, question_hash, question_text, "
            " answer_text, merkle_proof, model_profile_hash, conversation_hash, "
            " governance_policy_hash, schema_version, canonicalization_version, "
            " chunking_version, falsification_state, chain, audit_event_hash, "
            " created_at, hit_count) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'live', ?, ?, ?, 0)",
            (
                ckey,
                document_root,
                doc["document_uri"],
                qhash,
                question,
                answer_text,
                proof_blob,
                mhash,
                chash,
                ghash,
                SCHEMA_VERSION,
                CANONICALIZATION_VERSION,
                doc["chunking_version"],
                chain,
                event_hash,
                now,
            ),
        )

    return {
        "status": "cache_miss_then_written",
        "audit_mode": "STRICT",
        "cache_key": ckey,
        "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),
        },
    }