arborist/aborist/qa/query.py

"""Multi-source corpus Q&A.

Pose a question, the tree finds related cached docs, assembles them as context,
asks Hermes, caches the answer.

The flow:
    1. FTS5 search across all shards (chunks_fts can't be UNION'd in views,
       so we query each shard's index independently and merge by score).
    2. Pick top-K distinct documents within a character budget.
    3. Compute `context_root` = Merkle root over the sorted source
       document_roots — that's the "source" dimension of v9.8's 8-dim
       cache_key for this multi-source answer.
    4. Cache lookup; hit returns STRICT immediately.
    5. Miss calls Hermes via the OpenAI-compatible client.
    6. Persist record with merkle_proof = {context_root, sources: [...]}
       so a verifier can recompute context_root and check each source
       document_root against its shard.

Per-source proofs are not bundled here (the source roots themselves are
already content-addressed). A verifier asks the shards for any specific
chunk's proof on demand.
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass
from pathlib import Path

from aborist import (
    CANONICALIZATION_VERSION,
    CHUNKING_VERSION,
    SCHEMA_VERSION,
)
from aborist.merkle import MerkleTree
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.search import FTS5Backend
from aborist.store import (
    append_audit,
    connect,
    discover_shards,
    transaction,
)


DEFAULT_QUERY_POLICY = {
    "system_prompt": (
        "You are answering a question using ONLY the sources provided below. "
        "Each source is delimited by '=== Source: <URI> ===' headers. "
        "Cite which source supports each claim by quoting it inline. "
        "If the sources do not contain the answer, say 'I don't know based on "
        "the provided sources.' Do not speculate beyond what is written."
    ),
    "temperature": 0.1,
    "top_p": 1.0,
    "max_tokens": 768,
}


@dataclass
class _Hit:
    document_root: str
    document_uri: str
    title: str | None
    score: float
    shard_path: str
    chunk_idx: int


def _search_corpus(
    shards_dir: Path | None,
    single_db: Path | None,
    question: str,
    over_fetch: int,
) -> list[_Hit]:
    """FTS5 across shards (or single DB). Dedupe by document_root."""
    paths: list[Path]
    if shards_dir is not None:
        paths = discover_shards(shards_dir)
    elif single_db is not None:
        paths = [Path(single_db)]
    else:
        paths = []

    raw: list[tuple] = []
    for p in paths:
        conn = connect(p)
        try:
            backend = FTS5Backend(conn)
            hits = backend.search(question, limit=over_fetch)
            for h in hits:
                raw.append(
                    (
                        h.score,
                        h.document_root,
                        h.document_uri,
                        h.title,
                        h.chunk_idx,
                        str(p.resolve()),
                    )
                )
        finally:
            conn.close()

    raw.sort(key=lambda r: -r[0])
    seen: set[str] = set()
    out: list[_Hit] = []
    for score, root, uri, title, idx, sp in raw:
        if root in seen:
            continue
        seen.add(root)
        out.append(
            _Hit(
                document_root=root,
                document_uri=uri,
                title=title,
                score=score,
                shard_path=sp,
                chunk_idx=idx,
            )
        )
    return out


def _load_doc_text(shard_path: str, document_root: str) -> str | None:
    """Concatenate all hot chunks of a document. Returns None if cold or missing."""
    conn = connect(shard_path)
    try:
        rows = conn.execute(
            "SELECT content FROM chunks "
            "WHERE document_root = ? AND content IS NOT NULL "
            "ORDER BY idx ASC",
            (document_root,),
        ).fetchall()
    finally:
        conn.close()
    if not rows:
        return None
    return "\n\n".join(r["content"] for r in rows)


def _context_root(source_roots: list[str]) -> str:
    """Merkle root over sorted source document_roots — the v9.8 'source' dim
    for multi-source answers. Sorting makes the root deterministic regardless
    of search ranking order."""
    if not source_roots:
        return "00" * 32
    sorted_roots = sorted(source_roots)
    if len(sorted_roots) == 1:
        return sorted_roots[0]
    leaves = [bytes.fromhex(r) for r in sorted_roots]
    return MerkleTree.build(leaves).root.hex()


def query(
    *,
    question: str,
    qa_db: Path,
    chat_client: ChatClient,
    model_id: str,
    revision: str = "",
    quantization: str = "",
    shards_dir: Path | None = None,
    single_db: Path | None = None,
    top_k: int = 8,
    over_fetch: int = 32,
    max_context_chars: int = 60000,
    policy: dict | None = None,
    chain: str = "private",
) -> dict:
    """Answer `question` using the corpus. Cache to qa_db. Returns a result dict."""
    policy = policy or DEFAULT_QUERY_POLICY

    # 1. Search.
    hits = _search_corpus(shards_dir, single_db, question, over_fetch)
    if not hits:
        return {"status": "no_sources", "msg": "FTS5 search returned no hits"}

    # 2. Pull doc texts within budget.
    chosen: list[_Hit] = []
    context_parts: list[str] = []
    char_budget = max_context_chars
    for h in hits[:top_k]:
        text = _load_doc_text(h.shard_path, h.document_root)
        if not text:
            continue
        if len(text) > char_budget:
            text = text[:char_budget]
        context_parts.append(
            f"=== Source: {h.document_uri} ===\n{text}"
        )
        chosen.append(h)
        char_budget -= len(text)
        if char_budget <= 0:
            break

    if not chosen:
        return {"status": "no_sources", "msg": "top-k hits had cold or empty content"}

    context = "\n\n".join(context_parts)

    # 3. Build messages + hashes.
    messages = [
        {"role": "system", "content": policy["system_prompt"]},
        {
            "role": "user",
            "content": f"Sources:\n\n{context}\n\n---\n\nQuestion: {question}",
        },
    ]
    context_root = _context_root([h.document_root for h in chosen])
    qhash = question_hash(question)
    mhash = model_profile_hash(model_id, revision, quantization)
    chash = conversation_hash(messages)
    ghash = governance_policy_hash(policy)
    ckey = cache_key(
        context_root,
        qhash,
        mhash,
        chash,
        ghash,
        SCHEMA_VERSION,
        CANONICALIZATION_VERSION,
        CHUNKING_VERSION,
    )

    qa_conn = connect(qa_db)
    try:
        # 4. Cache lookup.
        cached = qa_conn.execute(
            "SELECT * FROM providence_cache "
            "WHERE cache_key = ? AND falsification_state = 'live'",
            (ckey,),
        ).fetchone()
        if cached is not None:
            now = int(time.time())
            with transaction(qa_conn):
                qa_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,
                "context_root": context_root,
                "answer_text": cached["answer_text"],
                "sources": json.loads(cached["merkle_proof"])["sources"],
            }

        # 5. Cache miss — call LLM.
        answer_text = chat_client.chat_completion(
            messages,
            model=model_id,
            temperature=policy["temperature"],
            max_tokens=policy["max_tokens"],
            top_p=policy.get("top_p", 1.0),
        )

        # 6. Persist record + audit event.
        proof_obj = {
            "context_root": context_root,
            "sources": [
                {
                    "document_root": h.document_root,
                    "document_uri": h.document_uri,
                    "title": h.title,
                    "score": h.score,
                    "chunk_idx": h.chunk_idx,
                    "shard": Path(h.shard_path).name,
                }
                for h in chosen
            ],
        }
        proof_blob = json.dumps(proof_obj, separators=(",", ":"))
        now = int(time.time())
        with transaction(qa_conn):
            event_hash = append_audit(
                qa_conn,
                event_type="providence_query",
                subject_root=ckey,
                body={
                    "context_root": context_root,
                    "n_sources": len(chosen),
                    "model_id": model_id,
                    "revision": revision,
                    "quantization": quantization,
                    "answer_chars": len(answer_text),
                    "context_chars": len(context),
                },
                ts=now,
            )
            qa_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,
                    context_root,
                    "corpus://multi-source",
                    qhash,
                    question,
                    answer_text,
                    proof_blob,
                    mhash,
                    chash,
                    ghash,
                    SCHEMA_VERSION,
                    CANONICALIZATION_VERSION,
                    CHUNKING_VERSION,
                    chain,
                    event_hash,
                    now,
                ),
            )
    finally:
        qa_conn.close()

    return {
        "status": "cache_miss_then_written",
        "audit_mode": "STRICT",
        "cache_key": ckey,
        "context_root": context_root,
        "answer_text": answer_text,
        "sources": proof_obj["sources"],
    }