brain-of-reese/scripts/load_test_kb.py

"""Build the controlled tool-calling test KB and snapshot it (one-off).

The fixture KB lives in ``tests/fixtures/agent_kb/`` — two source
directories (``deployments``, ``homelab``) with eight hand-written
markdown documents whose specifics (``rack7``, ``10.77.42.0/24``,
VLAN 130, port 18443, ntfy topic ``reese-uptime-7``, machine ID
``rbm-8842``, the ``17 2 * * *`` schedule, image
``ghcr.io/reese/obsidian-bor:2026.7.14``, port 18765, …) are not
guessable by any model. This script:

1. resets the app tables (one TRUNCATE — the dump's table set),
2. registers the two fixture directories as ``kind='local'``
   ``git_sources`` rows (so the source registry — and therefore
   ``ls``'s scope names — is self-contained and independent of the
   ``BOR_GIT_SOURCES`` env var),
3. imports the fixture documents through the real pipeline
   (``import_sources`` — real chunking + real ``embed``-model
   embeddings; this is the ONLY step that burns model calls, and only
   at build time),
4. stores the static KB overview + the sources-version row,
5. prints a **retrieval report** for every fixture-battery question
   (grounded or deflected, which documents would seed) — the battery
   must be all-grounded for the gate to exercise the tools,
6. snapshots the resulting database state into
   ``tests/fixtures/test_kb.dump.sql`` — a data-only SQL script
   (TRUNCATE + one multi-row ``INSERT`` per app table, generated
   in-process — the same file runs in psql or psycopg, in one
   transaction) — and **verifies the snapshot by restoring it and
   comparing a per-table checksum**.

Re-run it only when the fixture documents, the chunker, or the
embedding model change — everyday iterations restore the dump in
sub-second time (``scripts/restore_test_kb`` / the gate's
``--restore``), never re-embedding (see ``TOOL_CALLING_TESTING.md``).

Exit codes: **0** built + verified, **1** build/verification failure,
**2** precondition failure.
"""
from __future__ import annotations

import argparse
import asyncio
import json
import logging
import sys
import time
import uuid
from datetime import datetime
from pathlib import Path

from dotenv import load_dotenv
from sqlalchemy import select, text
from sqlalchemy.orm import Session

from app.api.chat import plan_turn
from app.config import get_settings
from app.db import SessionLocal, db_available
from app.models import (
    Chunk,
    DocDraft,
    Document,
    GitSource,
    KbOverview,
    QueryLog,
    SavedChat,
    SourcesMeta,
    SteeringNote,
)
from app.rag.importer import import_sources
from app.rag.llm import LLMClient
from app.rag.retriever import retrieve

logger = logging.getLogger("scripts.load_test_kb")

DEFAULT_KB_DIR = Path("tests/fixtures/agent_kb")
DEFAULT_DUMP = Path("tests/fixtures/test_kb.dump.sql")

#: The two fixture source directories (source name = directory basename,
#: the importer's rule). Alphabetical — the catalog order the derived
#: battery reads.
FIXTURE_SOURCES: tuple[str, ...] = ("deployments", "homelab")

#: The KB overview stored with the fixture (id=1). A plain outline of
#: the KB's basic categories — the ``<knowledge_base>`` prompt section
#: of every turn. Static on purpose: the dump must be deterministic and
#: the gate must not burn a ``lite`` call at restore time.
FIXTURE_KB_OVERVIEW: str = (
    "deployments: the lab Ansible inventory (host addresses and roles), "
    "the Obsidian BOR quadlet service definition, and the GitLab Runner "
    "CI setup. homelab: the rack7 Proxmox cluster networking (bridges, "
    "VLANs, DNS/DHCP), container notes (Uptime Kuma, Qwen 3.8 on "
    "llama.cpp), and the nightly restic backup configuration."
)

#: (table, model, explicit column list — the generated ``chunks.tsv``
#: tsvector column is excluded; Postgres recomputes it).
_TABLES: tuple[tuple[str, type, tuple[str, ...]], ...] = (
    ("documents", Document, ("id", "source", "path", "full_path", "title",
                             "content", "content_hash", "indexed_at", "summary")),
    ("chunks", Chunk, ("id", "document_id", "position", "content",
                       "embedding", "is_summary")),
    ("git_sources", GitSource, ("id", "url", "kind", "path", "added_at")),
    ("kb_overview", KbOverview, ("id", "content", "updated_at")),
    ("sources_meta", SourcesMeta, ("id", "version", "updated_at")),
    ("steering_notes", SteeringNote, ("id", "note", "created_at")),
    ("query_log", QueryLog, ("id", "question", "top_score", "fts_hits",
                             "chunk_hits", "deflected", "sources",
                             "latency_ms", "created_at")),
    ("saved_chats", SavedChat, ("id", "title", "messages", "share_token",
                                "sources_version", "created_at", "updated_at")),
    ("doc_drafts", DocDraft, ("id", "token", "title", "path", "body",
                              "status", "branch", "commit_sha", "created_at",
                              "updated_at")),
)


# --------------------------------------------------------------------------
# SQL serialization (the dump is plain multi-row INSERTs — the installed
# psycopg build exposes no COPY API, and a multi-statement script with
# inline ``COPY … FROM stdin`` data cannot be sent through any driver's
# simple-protocol execute. INSERT VALUES is the portable form: the same
# file runs in psql, psycopg, or anything else that speaks SQL, in one
# transaction. With ``standard_conforming_strings`` on (the Postgres
# default since 9.1), a string literal needs ONLY single-quote doubling —
# backslashes are literal and newlines may be real.
# --------------------------------------------------------------------------

def _sql_value(value: object) -> str:
    """One value as a SQL literal (``NULL`` for None)."""
    if value is None:
        return "NULL"
    if isinstance(value, bool):
        return "TRUE" if value else "FALSE"
    if isinstance(value, float):
        return repr(value)  # shortest round-trip double
    if isinstance(value, int):
        return str(value)
    if isinstance(value, uuid.UUID):
        return "'" + str(value) + "'"
    if isinstance(value, datetime):
        return "'" + value.isoformat(sep=" ") + "'"
    if isinstance(value, (list, tuple)) and value and isinstance(value[0], float):
        # A pgvector vector: the ``[v1, v2, …]`` text literal (pgvector
        # 0.7+ format; the older ``{…}`` form is rejected). Checked
        # before the JSONB branch — a JSONB array of dicts never has a
        # float first element.
        return "'" + "[" + ",".join(repr(v) for v in value) + "]" + "'"
    if isinstance(value, (dict, list)):
        # JSONB columns: the stored JSON text (Postgres re-parses it).
        text_ = json.dumps(value, ensure_ascii=False, separators=(",", ":"))
    else:
        text_ = str(value)
    return "'" + text_.replace("'", "''") + "'"


def _dump_table(db: Session, table: str, model: type, columns: tuple[str, ...]) -> str:
    """One multi-row ``INSERT INTO <table> (…) VALUES (…), …;`` statement
    (an empty table emits nothing — there is no row to write)."""
    rows = db.execute(select(model)).all()
    value_rows = [
        "(" + ", ".join(_sql_value(getattr(row[0], name)) for name in columns) + ")"
        for row in rows
    ]
    if not value_rows:
        return ""
    return (
        f"INSERT INTO public.{table} ({', '.join(columns)}) VALUES "
        + ",\n".join(value_rows)
        + ";\n"
    )


def _table_checksum(db: Session, table: str) -> str:
    """An order-independent content checksum for *table* (row::text,
    sorted aggregation) — the snapshot round-trip check."""
    return db.execute(
        text(
            "select md5(coalesce(string_agg(r, E'\\n' order by r), '')) "
            f"from (select t::text as r from public.{table} t) s"
        )
    ).scalar_one()


async def _retrieval_report(llm: LLMClient, battery: list[str]) -> int:
    """Print, per battery question, what the real path would do:
    grounded or deflected, and which documents would seed the context.
    Returns the number of deflected questions (a loud warning — the
    gate needs tools offered on its turns)."""
    settings = get_settings()
    deflected = 0
    print("\nretrieval report (the honesty gate per battery question):")
    with SessionLocal() as db:
        for number, question in enumerate(battery, start=1):
            vec = await llm.embed_one(question)
            chunks = retrieve(db, question, vec)
            plan = plan_turn(chunks, settings)
            seed = ", ".join(f"{d.source}/{d.path}" for d in plan.docs) or "—"
            if plan.deflected:
                deflected += 1
            print(
                f"  {number:02d}. {'DEFLECTED ' if plan.deflected else 'grounded '} "
                f"(best={plan.top_score:.3f} fts={plan.fts_hits}) "
                f"seed: {seed}\n      ← {question}"
            )
    return deflected


async def _build(kb_dir: Path, dump_path: Path) -> int:
    from scripts.agent_realmodel_check import FIXTURE_BATTERY

    started = time.monotonic()
    if not db_available():
        print(
            "load_test_kb: precondition failed — database unreachable; "
            "start Postgres with `podman compose up -d db`"
        )
        return 2

    source_dirs = [kb_dir / name for name in FIXTURE_SOURCES]
    missing = [str(p) for p in source_dirs if not p.is_dir()]
    if missing:
        print(f"load_test_kb: precondition failed — missing source dir(s): {missing}")
        return 2

    # 1. Reset the dump's table set (one statement — the inter-table FKs
    #    resolve within it).
    table_list = ", ".join(f"public.{table}" for table, _m, _c in _TABLES)
    with SessionLocal() as db:
        db.execute(text(f"TRUNCATE {table_list}"))
        for directory in source_dirs:
            absolute = str(directory.resolve())
            db.add(GitSource(url=absolute, kind="local", path=absolute))
        db.commit()
    logger.info("load_test_kb: tables reset; %d local source rows added", len(source_dirs))

    # 2. Import through the real pipeline (the only model-cost step).
    llm = LLMClient()
    summary = await import_sources([p.resolve() for p in source_dirs], llm)
    if summary.errors:
        print(f"load_test_kb: {summary.errors} file(s) failed to import — aborting")
        return 1
    if summary.added == 0:
        print("load_test_kb: no documents imported — aborting")
        return 1
    logger.info(
        "load_test_kb: imported added=%d chunks=%d embed_batches=%d",
        summary.added, summary.chunks, summary.embed_batches,
    )

    # 3. The static overview + sources version.
    with SessionLocal() as db:
        overview = db.get(KbOverview, 1) or KbOverview(id=1)
        overview.content = FIXTURE_KB_OVERVIEW
        meta = db.get(SourcesMeta, 1) or SourcesMeta(id=1)
        meta.version = 1
        db.add(overview)
        db.add(meta)
        db.commit()

    # 4. Retrieval report (all battery questions must stay grounded).
    n_deflected = await _retrieval_report(llm, list(FIXTURE_BATTERY))
    if n_deflected:
        print(
            f"\nload_test_kb: WARNING — {n_deflected} battery question(s) would "
            "DEFLECT in the real path (no tools offered). Adjust the fixture "
            "content (a lexical anchor for the question's words) or the "
            "question before running the gate."
        )

    # 5. Snapshot (data-only) + round-trip verification.
    with SessionLocal() as db:
        before = {table: _table_checksum(db, table) for table, _m, _c in _TABLES}
        parts = [
            "-- ============================================================",
            "-- Brain of Reese — controlled tool-calling test KB (fixture dump)",
            f"-- Generated by scripts/load_test_kb.py on "
            f"{datetime.now().astimezone().isoformat(timespec='seconds')}",
            "-- Data-only snapshot (the schema stays alembic-managed; the",
            "-- generated chunks.tsv column is recomputed on restore).",
            f"-- Sources: {', '.join(FIXTURE_SOURCES)} "
            f"({summary.added} documents, {summary.chunks} chunks).",
            "-- Restore (one transaction, sub-second):",
            "--   uv run python -m scripts.restore_test_kb",
            "--   psql \"$BOR_DATABASE_URL\" --single-transaction -f "
            "tests/fixtures/test_kb.dump.sql",
            "-- ============================================================",
            f"TRUNCATE {table_list};",
            "",
        ]
        for table, model, columns in _TABLES:
            parts.append(_dump_table(db, table, model, columns))
        script = "\n".join(parts)
    dump_path.parent.mkdir(parents=True, exist_ok=True)
    dump_path.write_text(script, encoding="utf-8")
    logger.info("load_test_kb: dump written: %s (%d KB)", dump_path, len(script) // 1024)

    # Round-trip: restore the dump over the (identical) state and compare
    # the per-table checksums — a serialization bug must fail the build.
    from scripts.restore_test_kb import restore_dump

    restore_dump(dump_path)  # RuntimeError on failure → the build fails
    with SessionLocal() as db:
        after = {table: _table_checksum(db, table) for table, _m, _c in _TABLES}
    mismatched = [t for t, c in before.items() if after.get(t) != c]
    if mismatched:
        print(f"load_test_kb: VERIFICATION FAILED — checksum mismatch: {mismatched}")
        return 1

    wall = time.monotonic() - started
    print(
        f"load_test_kb: ok — docs={summary.added} chunks={summary.chunks} "
        f"sources={len(FIXTURE_SOURCES)} dump={dump_path} "
        f"({dump_path.stat().st_size // 1024} KB, verified by round-trip) "
        f"in {wall:.1f}s"
    )
    return 0


def main(argv: list[str] | None = None) -> int:
    # CLI-only: pick up .env without side effects on import (the house
    # probe pattern, cf. scripts/llm_probe.py).
    load_dotenv()
    parser = argparse.ArgumentParser(
        description=(
            "Build the controlled tool-calling test KB from "
            "tests/fixtures/agent_kb (real embeddings, once) and snapshot "
            "it to tests/fixtures/test_kb.dump.sql (verified by "
            "round-trip). Exit 0 built+verified, 1 failure, 2 precondition."
        )
    )
    parser.add_argument(
        "--kb-dir", type=Path, default=DEFAULT_KB_DIR,
        help=f"the fixture KB root (default: {DEFAULT_KB_DIR})",
    )
    parser.add_argument(
        "--dump", type=Path, default=DEFAULT_DUMP,
        help=f"the dump file to write (default: {DEFAULT_DUMP})",
    )
    args = parser.parse_args(argv)
    logging.basicConfig(
        level=logging.INFO, format="%(levelname)s %(name)s: %(message)s"
    )
    return asyncio.run(_build(args.kb_dir, args.dump))


if __name__ == "__main__":
    sys.exit(main())