arborist/tests/test_query.py

"""Multi-source corpus Q&A: search → context → cache → Hermes.

Stub client only — no network. Validates:
- FTS5 finds the right docs across a small corpus
- context_root is deterministic (sorted source roots, then Merkle)
- cache hit on identical question returns same record without calling client
- different question -> different cache_key
"""

from __future__ import annotations

import json
from typing import Iterator

from aborist.document import Document
from aborist.ingest import ingest_source
from aborist.qa import query
from aborist.qa.client import StubClient
from aborist.source import Source
from aborist.store import connect


class FakeSource(Source):
    source_type = "test"

    def __init__(self, docs: list[Document]):
        self.docs = docs

    def iter_documents(self) -> Iterator[Document]:
        yield from self.docs


def _doc(uri: str, content: str) -> Document:
    return Document(uri=uri, content=content, source_type="test", title=uri.rsplit("/", 1)[-1])


# Three docs with distinguishable content so FTS5 can pick winners.
DOCS = [
    _doc(
        "test://anarchism",
        "Anarchism is a political philosophy that opposes the state. " * 12
        + "Mutual aid is central to anarchist theory. " * 8,
    ),
    _doc(
        "test://capitalism",
        "Capitalism is an economic system based on private ownership. " * 12
        + "Market exchange and capital accumulation drive growth. " * 8,
    ),
    _doc(
        "test://anarcho-capitalism",
        "Anarcho-capitalism combines anarchism's opposition to the state with capitalism's private property. " * 12
        + "Murray Rothbard developed many of its core ideas. " * 8,
    ),
]


def test_query_picks_relevant_docs_and_writes_record(tmp_path):
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(DOCS))
    finally:
        conn.close()

    client = StubClient(answer="anarcho-capitalism rests on private property without a state.")
    result = query(
        question="What is anarcho-capitalism?",
        qa_db=qa_db,
        chat_client=client,
        model_id="test-model",
        single_db=main_db,
        top_k=3,
    )

    assert result["status"] == "cache_miss_then_written"
    assert result["audit_mode"] == "STRICT"
    assert "anarcho-capitalism" in result["answer_text"]
    assert len(result["sources"]) >= 1
    assert any("anarcho-capitalism" in s["document_uri"] for s in result["sources"])

    # context_root is deterministic.
    sorted_roots = sorted(s["document_root"] for s in result["sources"])
    if len(sorted_roots) == 1:
        assert result["context_root"] == sorted_roots[0]
    # Repeat call with same question → cache hit, no LLM call.
    n_calls_before = len(client.calls)
    result2 = query(
        question="What is anarcho-capitalism?",
        qa_db=qa_db,
        chat_client=client,
        model_id="test-model",
        single_db=main_db,
        top_k=3,
    )
    assert result2["status"] == "cache_hit"
    assert result2["cache_key"] == result["cache_key"]
    assert len(client.calls) == n_calls_before  # no new call


def test_query_different_question_different_cache(tmp_path):
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(DOCS))
    finally:
        conn.close()

    client = StubClient(answer="answer text")
    r1 = query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=client,
        model_id="m",
        single_db=main_db,
    )
    r2 = query(
        question="What is capitalism?",
        qa_db=qa_db,
        chat_client=client,
        model_id="m",
        single_db=main_db,
    )
    assert r1["cache_key"] != r2["cache_key"]
    assert len(client.calls) == 2  # both missed cache, both called client


def test_query_no_sources_when_empty_corpus(tmp_path):
    main_db = tmp_path / "empty.db"
    qa_db = tmp_path / "qa.db"
    connect(main_db).close()  # creates schema, no docs
    result = query(
        question="What is anything?",
        qa_db=qa_db,
        chat_client=StubClient(),
        model_id="m",
        single_db=main_db,
    )
    assert result["status"] == "no_sources"


def test_query_persists_audit_event(tmp_path):
    main_db = tmp_path / "corpus.db"
    qa_db = tmp_path / "qa.db"
    conn = connect(main_db)
    try:
        ingest_source(conn, FakeSource(DOCS))
    finally:
        conn.close()
    query(
        question="What is anarchism?",
        qa_db=qa_db,
        chat_client=StubClient(answer="X"),
        model_id="m",
        single_db=main_db,
    )
    # qa.db should now have one providence_query event in its audit chain.
    qc = connect(qa_db)
    try:
        events = qc.execute(
            "SELECT event_type FROM audit_events ORDER BY seq"
        ).fetchall()
    finally:
        qc.close()
    types = [e["event_type"] for e in events]
    assert "providence_query" in types