tai/tests/rag.py

#!/usr/bin/env python3
"""tests/rag.py — integration tests for tai's chunked FTS5/BM25 RAG system.

Built on tai_test.py (see tests/README.md for the conventions this file
follows). Two test cases, deliberately different in how deterministic they
are:

- InteractiveReplSmokeTest drives a persistent REPL over a real pty with
  nothing but REPL slash commands (/skills, /search). Slash commands dispatch
  tools directly, with no LLM call involved, so this test is fully
  deterministic and fast — it exists to prove the pty transport, the
  prompt-detection turn boundary, and the REPL plumbing all actually work,
  independent of model behavior.
- RagIndexAndSearchTest is the real ask: index a sample of this project's own
  research/ archive into a fresh, isolated vault, then verify it comes back
  through search. Indexing (asking the agent to call record_save) goes
  through the LLM, which is inherently probabilistic — a free, rotating
  backend (see README.md's "Backends" section) can pick a different model
  per request. This test accepts that honestly: it reports how many of the
  sampled files were actually indexed rather than demanding all-or-nothing,
  and only hard-fails if indexing came back empty or the deterministic
  search step (itself a REPL slash command, no LLM) doesn't find anything
  that was actually indexed.
"""
import difflib
import os
import re
import shutil
import sqlite3
import sys
import tempfile
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from tai_test import TestCase, TestAssertionFailed, main_cli  # noqa: E402

ROOT = Path(__file__).resolve().parent.parent
RESEARCH_DIR = ROOT / "research"
REPL_PROMPT_RE = re.compile(r"tai\[[^\]]*\]\s*[›>]")
SAMPLE_SIZE = int(os.environ.get("RAG_TEST_SAMPLE_SIZE", "5"))
SOURCE_URL_RE = re.compile(r"\|\s*Source URL\s*\|\s*(\S+)\s*\|")


def pick_sample_files(limit=SAMPLE_SIZE):
    # Only source dossiers/citation cards: each has exactly one "Source URL"
    # metadata field, which is what gives us a per-file, globally-unique
    # fingerprint. Conclusions cite several URLs each, so they don't fit
    # this scheme and are excluded.
    files = sorted(
        p for p in RESEARCH_DIR.rglob("*.md")
        if p.name not in ("README.md", "INDEX.md") and "conclusions" not in p.parts
    )
    return files[:limit]


class InteractiveReplSmokeTest(TestCase):
    """Deterministic: pure REPL slash commands, no LLM call, proves the pty transport."""

    name = "repl-pty-smoke"
    mode = "pty"
    timeout = 30.0

    async def setup(self):
        self.tmp_home = tempfile.mkdtemp(prefix="tai-test-home-")
        self.env_overrides = {"TAI_HOME": self.tmp_home, "TAI_PASSPHRASE": ""}

    async def run(self):
        proc = await self.spawn(extra_args=["--profile", "smoketest"])
        await proc.wait_for_pattern(REPL_PROMPT_RE, timeout=15.0)
        await proc.send_line("/skills")
        skills_reply = await proc.wait_for_pattern(REPL_PROMPT_RE, timeout=15.0)
        await proc.send_line("/search nonexistent-term-xyz")
        search_reply = await proc.wait_for_pattern(REPL_PROMPT_RE, timeout=15.0)
        await proc.send_line("/quit")
        await proc.wait(timeout=10.0)
        return {"skills_reply": skills_reply, "search_reply": search_reply}

    async def verify(self, result):
        if "skills" not in result["skills_reply"] and "no skills" not in result["skills_reply"]:
            raise TestAssertionFailed("/skills produced no recognizable output: %r" % result["skills_reply"][:300])
        if "no matches" not in result["search_reply"] and "error" in result["search_reply"].lower():
            raise TestAssertionFailed("/search errored instead of reporting no matches: %r" % result["search_reply"][:300])

    async def teardown(self):
        shutil.rmtree(self.tmp_home, ignore_errors=True)


class RagIndexAndSearchTest(TestCase):
    """The real ask, proven with per-file traceability, not just 'something came back'.

    For each sampled file this test:
      1. extracts that file's Source URL as a fingerprint, and confirms it's
         not shared by any other file in the sample (a non-unique fingerprint
         would make a search hit ambiguous as "proof");
      2. asks the agent to record_save it, exactly as before;
      3. reads the saved record straight out of the vault DB and diffs it
         against what was actually sent (not the whole file — what we asked
         the model to save verbatim), catching silent paraphrasing/truncation;
      4. confirms the fingerprint survived chunking, scoped to that record's
         own chunks specifically (not "some chunk somewhere");
      5. runs a /search for that one file's fingerprint and requires the
         reply to name that exact record's mem: id.

    Every step is reported per file (see verify()'s failure message and the
    log), so a partial result is auditable rather than a single opaque
    pass/fail.
    """

    name = "rag-index-and-search"
    mode = "pipe"
    # Per file: one save round trip (up to SAVE_TIMEOUT, observed free-tier backends taking
    # 30-60s even on success) + one search round trip (no LLM call, local FTS, cheap) + process
    # wait/spawn overhead. This must scale with SAMPLE_SIZE (env-configurable, default 5) -- a
    # flat constant sized for an earlier, smaller default silently became too tight and made a
    # slow-but-successful free-tier response get killed as a false "timeout" failure.
    SAVE_TIMEOUT = 180.0
    timeout = SAMPLE_SIZE * (SAVE_TIMEOUT + 10.0 + 15.0 + 15.0 + 10.0) + 60.0

    async def setup(self):
        self.tmp_home = tempfile.mkdtemp(prefix="tai-test-home-")
        self.env_overrides = {"TAI_HOME": self.tmp_home, "TAI_PASSPHRASE": ""}
        self.files = pick_sample_files()
        if not self.files:
            raise TestAssertionFailed("no documents found under research/ to index — did research/README.md's archive get built?")

    async def run(self):
        fingerprints = self._extract_and_verify_unique_fingerprints(self.files)

        per_file = []
        for path, fingerprint in fingerprints:
            sent_content = path.read_text(encoding="utf-8", errors="replace")[:3000]
            title = path.relative_to(RESEARCH_DIR).as_posix()
            prompt = (
                "Call the record_save tool right now with kind=\"research\", "
                "title=\"%s\", and content set to exactly the text below, verbatim. "
                "Do not summarize it, do not explain yourself, just call the tool once.\n\n%s"
                % (title, sent_content)
            )
            proc = await self.spawn(extra_args=["--yes", "--profile", "ragtest", prompt])
            await proc.drain_until_eof(timeout=self.SAVE_TIMEOUT)
            await proc.wait(timeout=10.0)

            record = self._find_record_by_title(title)
            fidelity = self._fidelity_ratio(sent_content, record["content"]) if record else 0.0
            chunk_has_fingerprint = self._record_chunk_contains(record["id"], fingerprint) if record else False

            per_file.append({
                "title": title,
                "fingerprint": fingerprint,
                "saved": record is not None,
                "record_id": record["id"] if record else None,
                "fidelity": fidelity,
                "chunk_has_fingerprint": chunk_has_fingerprint,
            })

        search_results = []
        for entry in per_file:
            if not entry["saved"]:
                continue
            proc = await self.spawn(extra_args=["--profile", "ragtest"], mode="pty")
            await proc.wait_for_pattern(REPL_PROMPT_RE, timeout=15.0)
            await proc.send_line("/search %s" % entry["fingerprint"])
            reply = await proc.wait_for_pattern(REPL_PROMPT_RE, timeout=15.0)
            await proc.send_line("/quit")
            await proc.wait(timeout=10.0)
            found_own_record = entry["record_id"] in reply
            search_results.append({"title": entry["title"], "reply": reply, "found_own_record": found_own_record})
            entry["search_found_own_record"] = found_own_record

        return {"attempted": len(self.files), "per_file": per_file}

    def _extract_and_verify_unique_fingerprints(self, files):
        pairs = []
        seen = {}
        for path in files:
            text = path.read_text(encoding="utf-8", errors="replace")
            match = SOURCE_URL_RE.search(text)
            if not match:
                raise TestAssertionFailed("%s has no 'Source URL' metadata row — pick_sample_files() should have excluded it" % path)
            fingerprint = match.group(1)
            if fingerprint in seen:
                raise TestAssertionFailed(
                    "fingerprint %r is not unique: shared by %s and %s — cannot use it as proof of traceability"
                    % (fingerprint, seen[fingerprint], path)
                )
            seen[fingerprint] = path
            pairs.append((path, fingerprint))
        return pairs

    def _connect(self):
        return sqlite3.connect(str(Path(self.tmp_home) / "memory.db"))

    def _find_record_by_title(self, title):
        conn = self._connect()
        try:
            row = conn.execute("SELECT id, content FROM records WHERE kind = 'research' AND title = ?", (title,)).fetchone()
        finally:
            conn.close()
        return {"id": row[0], "content": row[1]} if row else None

    def _record_chunk_contains(self, record_id, fingerprint):
        conn = self._connect()
        try:
            rows = conn.execute("SELECT text FROM chunks WHERE parent_kind = 'record' AND parent_id = ?", (record_id,)).fetchall()
        finally:
            conn.close()
        return any(fingerprint in row[0] for row in rows)

    @staticmethod
    def _fidelity_ratio(sent, saved):
        if not saved:
            return 0.0
        return difflib.SequenceMatcher(None, sent, saved).ratio()

    async def verify(self, result):
        per_file = result["per_file"]
        report_lines = [
            "%-60s saved=%-5s fidelity=%.2f chunked=%-5s found_by_search=%-5s"
            % (e["title"], e["saved"], e["fidelity"], e["chunk_has_fingerprint"], e.get("search_found_own_record"))
            for e in per_file
        ]
        report = "per-file breakdown:\n  " + "\n  ".join(report_lines)

        saved_count = sum(1 for e in per_file if e["saved"])
        if saved_count == 0:
            raise TestAssertionFailed(
                "indexed 0 of %d sampled files (checked the vault's records table directly) — this may be "
                "free-backend flakiness, not a RAG bug; re-run before treating as a regression\n%s"
                % (result["attempted"], report)
            )

        fully_proven = [
            e for e in per_file
            if e["saved"] and e["fidelity"] >= 0.90 and e["chunk_has_fingerprint"] and e.get("search_found_own_record")
        ]
        if not fully_proven:
            raise TestAssertionFailed(
                "saved %d/%d file(s), but none passed the full chain (fidelity>=0.90 AND fingerprint survived "
                "chunking AND /search for that file's own fingerprint named that exact record) — no file is "
                "provably traceable end-to-end\n%s" % (saved_count, result["attempted"], report)
            )

        low_fidelity = [e for e in per_file if e["saved"] and e["fidelity"] < 0.90]
        if low_fidelity:
            raise TestAssertionFailed(
                "%d/%d saved file(s) had content fidelity below 0.90 against what was actually sent — the model "
                "paraphrased or truncated instead of saving verbatim: %s\n%s"
                % (len(low_fidelity), saved_count, [e["title"] for e in low_fidelity], report)
            )

        not_chunked = [e for e in per_file if e["saved"] and not e["chunk_has_fingerprint"]]
        if not_chunked:
            raise TestAssertionFailed(
                "%d/%d saved file(s) lost their fingerprint during chunking (present in the record but not in any "
                "of that record's own chunks): %s\n%s" % (len(not_chunked), saved_count, [e["title"] for e in not_chunked], report)
            )

        not_found = [e for e in per_file if e["saved"] and not e.get("search_found_own_record")]
        if not_found:
            raise TestAssertionFailed(
                "%d/%d saved file(s) did not come back when searching for their own unique fingerprint: %s\n%s"
                % (len(not_found), saved_count, [e["title"] for e in not_found], report)
            )

        print(report)

    async def teardown(self):
        shutil.rmtree(self.tmp_home, ignore_errors=True)


ALL_CASES = [InteractiveReplSmokeTest(), RagIndexAndSearchTest()]

if __name__ == "__main__":
    sys.exit(main_cli(ALL_CASES, default_log_dir=str(Path(__file__).resolve().parent / "logs")))