scripts/test_agent_chat.py

"""End-to-end test of Agent Chat: tools + RAG, with and without OpenRouter.

Covers:
  1. LanceDB index is built from the live HF dataset
  2. Tool executors work standalone (no LLM)
  3. Agentic flow (OpenRouter) — LLM calls search_records / get_records_by_date_range
  4. Fallback flow (OpenRouter disabled) — one-shot RAG via HF/no-op

Run:  python -m scripts.test_agent_chat
"""

from __future__ import annotations

import json
import logging
import os
import sys
from pathlib import Path

_PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_PROJECT_ROOT))

from dotenv import load_dotenv

load_dotenv(_PROJECT_ROOT / ".env")

logging.basicConfig(level=logging.WARNING, format="%(levelname)s: %(message)s")

CHIKU_PET_ID = "ec8ed3df-3b7b-4013-8d6d-c76832e1ec88"

SEP = "=" * 70


def _ok(cond: bool, label: str, detail: str = "") -> None:
    tag = "PASS" if cond else "FAIL"
    print(f"  [{tag}] {label}" + (f"  — {detail}" if detail else ""))


def step_1_build_index() -> None:
    print(SEP)
    print("STEP 1: pull latest HF dataset + rebuild LanceDB index")
    print(SEP)
    from guardiantails.services.hf_sync import download_backup
    from guardiantails.models import init_db, get_db
    from guardiantails.services.lance_service import rebuild_from_sqlite

    download_backup()
    init_db()
    with get_db() as conn:
        result = rebuild_from_sqlite(conn, pet_id=None)
    _ok(result.get("status") == "ok", "index rebuild", str(result))
    assert result.get("documents", 0) > 0, "no documents indexed — expected >0"
    print()


def step_2_tool_executors() -> None:
    print(SEP)
    print("STEP 2: tool executors (no LLM)")
    print(SEP)
    from guardiantails.core.rag_chat import (
        _tool_search_records,
        _tool_date_range,
    )

    # 2a. search_records
    r = json.loads(_tool_search_records(
        CHIKU_PET_ID,
        {"query": "seizure", "record_type": "seizure", "top_k": 5},
    ))
    _ok(isinstance(r, list) and len(r) > 0, "search_records returns seizures", f"got {len(r)}")
    if r:
        first = r[0]
        _ok(first.get("record_type") == "seizure", "first result is a seizure")

    # 2b. date-range: doses in 48h before the Apr 10 seizure
    r = json.loads(_tool_date_range(
        CHIKU_PET_ID,
        {"start_date": "2026-04-08", "end_date": "2026-04-10", "record_type": "dose"},
    ))
    doses_in_window = [x for x in r if x.get("record_type") == "dose"]
    missed = [x for x in doses_in_window if "Missed" in (x.get("summary") or "")]
    _ok(len(doses_in_window) > 0, "date-range returns doses",
        f"total={len(doses_in_window)}, missed={len(missed)}")

    # 2c. date-range unfiltered — should cover multiple types
    r = json.loads(_tool_date_range(
        CHIKU_PET_ID,
        {"start_date": "2026-04-08", "end_date": "2026-04-10"},
    ))
    types_found = {x.get("record_type") for x in r}
    _ok(len(types_found) >= 2, "date-range spans multiple record types",
        f"types={types_found}")
    print()


def step_3_agentic_flow() -> None:
    print(SEP)
    print("STEP 3: agentic flow via OpenRouter (should emit tool_calls)")
    print(SEP)
    if not os.getenv("OPENROUTER_API_KEY"):
        print("  [SKIP] OPENROUTER_API_KEY not set")
        return

    from guardiantails.models import get_db
    from guardiantails.core.rag_chat import rag_query

    question = (
        "Were there any missed or delayed medication doses in the 48 hours before "
        "Chiku's April 10 2026 seizure? Use the tools to check."
    )
    with get_db() as conn:
        out = rag_query(conn, CHIKU_PET_ID, question)

    trace = out.get("trace", [])
    tool_names = [t.get("name") for t in trace]
    _ok(len(trace) > 0, "LLM invoked at least one tool", f"calls={tool_names}")
    _ok(len(out.get("response", "")) > 50, "got a substantive response",
        f"len={len(out.get('response',''))}")
    _ok(len(out.get("sources", [])) > 0, "sources collected from tool calls",
        f"count={len(out.get('sources', []))}")
    print()
    print("  --- RESPONSE (truncated) ---")
    print("  " + (out.get("response") or "")[:600].replace("\n", "\n  "))
    print()


def step_4_fallback_flow() -> None:
    print(SEP)
    print("STEP 4: fallback (OpenRouter disabled) — one-shot RAG")
    print(SEP)

    # Temporarily strip the key and reload settings
    saved = os.environ.pop("OPENROUTER_API_KEY", None)
    try:
        from guardiantails.models import get_db
        from guardiantails.core.rag_chat import rag_query

        question = "Summarize Chiku's seizure history and any patterns."
        with get_db() as conn:
            out = rag_query(conn, CHIKU_PET_ID, question)

        _ok(isinstance(out.get("response"), str),
            "fallback returns a response string",
            f"len={len(out.get('response','') or '')}")
        _ok(len(out.get("sources", [])) > 0,
            "fallback retrieved sources via one-shot RAG",
            f"count={len(out.get('sources', []))}")
    finally:
        if saved is not None:
            os.environ["OPENROUTER_API_KEY"] = saved
    print()


def main() -> None:
    step_1_build_index()
    step_2_tool_executors()
    step_3_agentic_flow()
    step_4_fallback_flow()
    print(SEP)
    print("END-TO-END TEST COMPLETE")
    print(SEP)


if __name__ == "__main__":
    main()