src/backend/test_rag_quality.py

"""
RAG Retrieval Quality Test

Directly tests the ChromaDB vector search to validate that the expanded
clinical guidelines corpus returns relevant results for diverse clinical queries.

This bypasses the full agent pipeline and tests the retrieval layer in isolation.

Usage:
    cd src/backend
    python test_rag_quality.py                   # Run all queries
    python test_rag_quality.py --rebuild          # Delete ChromaDB and rebuild from scratch
    python test_rag_quality.py --stats            # Show collection statistics only
    python test_rag_quality.py --query "chest pain evaluation"  # Test a single query
"""
import asyncio
import argparse
import json
import shutil
import sys
import os
from pathlib import Path

# Ensure the app package is importable
sys.path.insert(0, str(Path(__file__).parent))

os.environ.setdefault("MEDGEMMA_API_KEY", "dummy")  # Prevent settings validation error


# ─────────────────────────────────────────────────
# Test Queries: (query, expected_specialty, expected_title_substring)
# Each query simulates what the orchestrator would send to the RAG tool.
# ─────────────────────────────────────────────────

RAG_TEST_QUERIES = [
    # Cardiology
    {
        "query": "Acute chest pain evaluation and troponin testing guidelines",
        "expected_specialties": ["Cardiology"],
        "expected_title_keywords": ["chest pain", "ACS", "STEMI"],
        "min_relevance": 0.3,
    },
    {
        "query": "Heart failure management with reduced ejection fraction HFrEF",
        "expected_specialties": ["Cardiology"],
        "expected_title_keywords": ["heart failure"],
        "min_relevance": 0.3,
    },
    {
        "query": "Atrial fibrillation anticoagulation and rate control",
        "expected_specialties": ["Cardiology"],
        "expected_title_keywords": ["atrial fibrillation", "AFib"],
        "min_relevance": 0.3,
    },
    {
        "query": "Management of acute pulmonary embolism with hemodynamic instability",
        "expected_specialties": ["Cardiology", "Pulmonology", "Hematology"],
        "expected_title_keywords": ["pulmonary embolism", "PE", "VTE"],
        "min_relevance": 0.25,
    },
    # Emergency Medicine
    {
        "query": "Acute ischemic stroke tPA thrombolysis eligibility criteria",
        "expected_specialties": ["Emergency Medicine", "Neurology"],
        "expected_title_keywords": ["stroke", "CVA"],
        "min_relevance": 0.3,
    },
    {
        "query": "Sepsis hour-1 bundle treatment with IV fluids and antibiotics",
        "expected_specialties": ["Emergency Medicine"],
        "expected_title_keywords": ["sepsis"],
        "min_relevance": 0.3,
    },
    {
        "query": "Anaphylaxis emergency treatment with epinephrine",
        "expected_specialties": ["Emergency Medicine"],
        "expected_title_keywords": ["anaphylaxis"],
        "min_relevance": 0.3,
    },
    {
        "query": "Trauma ATLS assessment primary survey hemorrhagic shock",
        "expected_specialties": ["Emergency Medicine"],
        "expected_title_keywords": ["trauma"],
        "min_relevance": 0.25,
    },
    {
        "query": "Seizure management status epilepticus benzodiazepine protocol",
        "expected_specialties": ["Emergency Medicine", "Neurology"],
        "expected_title_keywords": ["seizure"],
        "min_relevance": 0.25,
    },
    # Endocrinology
    {
        "query": "Diabetic ketoacidosis DKA insulin drip and fluid management",
        "expected_specialties": ["Endocrinology"],
        "expected_title_keywords": ["DKA", "diabetic ketoacidosis"],
        "min_relevance": 0.3,
    },
    {
        "query": "Type 2 diabetes management metformin A1C targets",
        "expected_specialties": ["Endocrinology"],
        "expected_title_keywords": ["diabetes", "DM"],
        "min_relevance": 0.3,
    },
    {
        "query": "Thyroid disease hyperthyroidism Graves disease treatment",
        "expected_specialties": ["Endocrinology"],
        "expected_title_keywords": ["thyroid"],
        "min_relevance": 0.3,
    },
    # Pulmonology
    {
        "query": "COPD exacerbation treatment bronchodilators steroids antibiotics",
        "expected_specialties": ["Pulmonology"],
        "expected_title_keywords": ["COPD"],
        "min_relevance": 0.3,
    },
    {
        "query": "Acute asthma exacerbation treatment albuterol magnesium",
        "expected_specialties": ["Pulmonology", "Pediatrics"],
        "expected_title_keywords": ["asthma"],
        "min_relevance": 0.3,
    },
    {
        "query": "Community acquired pneumonia antibiotic selection CURB-65",
        "expected_specialties": ["Pulmonology", "Infectious Disease"],
        "expected_title_keywords": ["pneumonia"],
        "min_relevance": 0.3,
    },
    # Gastroenterology
    {
        "query": "Upper GI bleeding management endoscopy PPI transfusion",
        "expected_specialties": ["Gastroenterology"],
        "expected_title_keywords": ["GI bleed", "upper GI"],
        "min_relevance": 0.3,
    },
    {
        "query": "Acute pancreatitis management fluid resuscitation pain control",
        "expected_specialties": ["Gastroenterology"],
        "expected_title_keywords": ["pancreatitis"],
        "min_relevance": 0.3,
    },
    # Neurology
    {
        "query": "Epilepsy seizure medication selection antiepileptic drugs",
        "expected_specialties": ["Neurology"],
        "expected_title_keywords": ["epilepsy"],
        "min_relevance": 0.3,
    },
    {
        "query": "Bacterial meningitis empiric antibiotics lumbar puncture",
        "expected_specialties": ["Neurology", "Infectious Disease"],
        "expected_title_keywords": ["meningitis"],
        "min_relevance": 0.3,
    },
    # Psychiatry
    {
        "query": "Suicide risk assessment safety planning lethal means counseling",
        "expected_specialties": ["Psychiatry"],
        "expected_title_keywords": ["suicide", "suicid"],
        "min_relevance": 0.3,
    },
    {
        "query": "Major depressive disorder SSRI treatment algorithm",
        "expected_specialties": ["Psychiatry"],
        "expected_title_keywords": ["depression", "depressive"],
        "min_relevance": 0.3,
    },
    # Pediatrics
    {
        "query": "Neonatal fever sepsis workup guidelines for infants under 60 days",
        "expected_specialties": ["Pediatrics"],
        "expected_title_keywords": ["fever", "neonate", "neonatal"],
        "min_relevance": 0.25,
    },
    {
        "query": "Pediatric dehydration oral rehydration IV fluid bolus",
        "expected_specialties": ["Pediatrics"],
        "expected_title_keywords": ["dehydration"],
        "min_relevance": 0.25,
    },
    # Nephrology
    {
        "query": "Hyperkalemia emergency management calcium insulin kayexalate dialysis",
        "expected_specialties": ["Nephrology", "Emergency Medicine"],
        "expected_title_keywords": ["hyperkalemia"],
        "min_relevance": 0.25,
    },
    {
        "query": "Acute kidney injury management and staging KDIGO",
        "expected_specialties": ["Nephrology"],
        "expected_title_keywords": ["AKI", "kidney injury"],
        "min_relevance": 0.25,
    },
    # Hematology
    {
        "query": "Venous thromboembolism DVT PE anticoagulation treatment duration",
        "expected_specialties": ["Hematology", "Cardiology"],
        "expected_title_keywords": ["VTE", "thromboembolism"],
        "min_relevance": 0.25,
    },
    # Infectious Disease
    {
        "query": "HIV antiretroviral therapy guidelines initial regimen",
        "expected_specialties": ["Infectious Disease"],
        "expected_title_keywords": ["HIV"],
        "min_relevance": 0.3,
    },
    {
        "query": "Urinary tract infection treatment pyelonephritis uncomplicated cystitis",
        "expected_specialties": ["Infectious Disease"],
        "expected_title_keywords": ["UTI", "urinary tract"],
        "min_relevance": 0.25,
    },
    # OB/GYN
    {
        "query": "Preeclampsia management magnesium sulfate antihypertensives",
        "expected_specialties": ["OB/GYN"],
        "expected_title_keywords": ["preeclampsia", "hypertensive"],
        "min_relevance": 0.3,
    },
    # Rheumatology
    {
        "query": "Acute gout treatment colchicine NSAIDs corticosteroids",
        "expected_specialties": ["Rheumatology"],
        "expected_title_keywords": ["gout"],
        "min_relevance": 0.3,
    },
]


async def rebuild_chroma(persist_dir: str):
    """Delete and recreate the ChromaDB collection."""
    p = Path(persist_dir)
    if p.exists():
        shutil.rmtree(p)
        print(f"  Deleted ChromaDB directory: {p}")
    else:
        print(f"  ChromaDB directory does not exist: {p}")

    # Re-init by creating a new tool instance and triggering load
    from app.tools.guideline_retrieval import GuidelineRetrievalTool
    tool = GuidelineRetrievalTool()
    await tool._ensure_initialized()
    assert tool._collection is not None, "Collection failed to initialize"
    count = tool._collection.count()
    print(f"  Rebuilt collection with {count} guidelines")
    return tool


async def show_stats(persist_dir: str):
    """Show ChromaDB collection statistics."""
    from app.tools.guideline_retrieval import GuidelineRetrievalTool
    tool = GuidelineRetrievalTool()
    await tool._ensure_initialized()
    assert tool._collection is not None, "Collection failed to initialize"

    count = tool._collection.count()
    print(f"\n  Collection: clinical_guidelines")
    print(f"  Documents: {count}")
    print(f"  Persist dir: {persist_dir}")

    if count > 0:
        # Get all metadata to show specialties
        all_data = tool._collection.get(include=["metadatas"])
        specialties = {}
        for meta in all_data["metadatas"]:
            spec = meta.get("specialty", "Unknown")
            specialties[spec] = specialties.get(spec, 0) + 1

        print(f"\n  Guidelines by specialty:")
        for spec, cnt in sorted(specialties.items()):
            print(f"    {spec:30s} {cnt}")

    return tool


async def test_single_query(tool, query_text: str, n_results: int = 5):
    """Test a single query and display results."""
    result = await tool.run(query_text, n_results=n_results)
    print(f"\n  Query: \"{query_text}\"")
    print(f"  Results: {len(result.excerpts)}")
    for i, exc in enumerate(result.excerpts):
        print(f"\n    [{i+1}] {exc.title}")
        print(f"        Source: {exc.source}")
        print(f"        Relevance: {exc.relevance_score:.4f}")
        print(f"        Excerpt: {exc.excerpt[:150]}...")


async def run_quality_tests(tool, test_queries):
    """Run all quality test queries and score results."""
    results = []

    for tq in test_queries:
        query = tq["query"]
        expected_specs = tq["expected_specialties"]
        expected_keywords = tq["expected_title_keywords"]
        min_rel = tq["min_relevance"]

        result = await tool.run(query, n_results=5)

        # Get top result info
        top_excerpt = result.excerpts[0] if result.excerpts else None
        top_title = top_excerpt.title if top_excerpt else "N/A"
        top_relevance = top_excerpt.relevance_score if top_excerpt else 0
        top_source = top_excerpt.source if top_excerpt else "N/A"

        # Check if any of the top-3 results match expected specialty
        specialty_match = False
        keyword_match = False
        matched_result_idx = -1

        for idx, exc in enumerate(result.excerpts[:3]):
            # Check source text or title for specialty/keyword matches
            title_lower = exc.title.lower()
            source_lower = exc.source.lower()
            combined = title_lower + " " + source_lower + " " + exc.excerpt.lower()

            for kw in expected_keywords:
                if kw.lower() in combined:
                    keyword_match = True
                    if matched_result_idx == -1:
                        matched_result_idx = idx
                    break

        # Relevance check
        relevance_ok = top_relevance >= min_rel

        # Overall pass: keyword match in top-3 AND minimum relevance
        passed = keyword_match and relevance_ok

        test_result = {
            "query": query[:60] + ("..." if len(query) > 60 else ""),
            "expected_specialties": expected_specs,
            "expected_keywords": expected_keywords,
            "top_title": top_title,
            "top_relevance": top_relevance,
            "keyword_match": keyword_match,
            "keyword_match_position": matched_result_idx + 1 if matched_result_idx >= 0 else 0,
            "relevance_ok": relevance_ok,
            "passed": passed,
            "all_titles": [e.title for e in result.excerpts[:5]],
            "all_relevances": [e.relevance_score for e in result.excerpts[:5]],
        }
        results.append(test_result)

    return results


async def main():
    parser = argparse.ArgumentParser(description="RAG Retrieval Quality Test")
    parser.add_argument("--rebuild", action="store_true", help="Rebuild ChromaDB from scratch")
    parser.add_argument("--stats", action="store_true", help="Show collection statistics only")
    parser.add_argument("--query", help="Test a single query")
    parser.add_argument("--verbose", action="store_true", help="Show detailed results for each query")
    args = parser.parse_args()

    from app.config import settings
    persist_dir = settings.chroma_persist_dir

    print(f"\n{'='*70}")
    print(f"  RAG Retrieval Quality Test")
    print(f"  Persist dir: {persist_dir}")
    print(f"  Embedding model: {settings.embedding_model}")
    print(f"{'='*70}")

    if args.rebuild:
        tool = await rebuild_chroma(persist_dir)
    elif args.stats:
        await show_stats(persist_dir)
        return
    else:
        from app.tools.guideline_retrieval import GuidelineRetrievalTool
        tool = GuidelineRetrievalTool()
        await tool._ensure_initialized()
        assert tool._collection is not None, "Collection failed to initialize"
        count = tool._collection.count()
        print(f"\n  Collection has {count} documents")
        if count == 0:
            print("  ⚠ Collection is empty! Run with --rebuild to load guidelines.")
            return

    if args.query:
        await test_single_query(tool, args.query)
        return

    # Run all quality tests
    print(f"\n  Running {len(RAG_TEST_QUERIES)} retrieval quality tests...\n")
    results = await run_quality_tests(tool, RAG_TEST_QUERIES)

    # Display results
    passed_count = 0
    for r in results:
        icon = "✓" if r["passed"] else "✗"
        pos = f"@{r['keyword_match_position']}" if r["keyword_match"] else "  "
        rel = f"{r['top_relevance']:.3f}"
        print(f"  {icon} [{rel}] {pos:>3} {r['query']}")
        if not r["passed"] or args.verbose:
            print(f"       → Top: {r['top_title']}")
            if not r["keyword_match"]:
                print(f"       ✗ Expected keywords not found in top-3: {r['expected_keywords']}")
            if not r["relevance_ok"]:
                print(f"       ✗ Relevance {r['top_relevance']:.3f} below threshold")
            if args.verbose:
                for i, (t, s) in enumerate(zip(r["all_titles"], r["all_relevances"])):
                    print(f"         {i+1}. [{s:.3f}] {t}")
        if r["passed"]:
            passed_count += 1

    # Summary
    total = len(results)
    pct = (passed_count / total * 100) if total else 0
    print(f"\n{'='*70}")
    print(f"  RESULTS: {passed_count}/{total} passed ({pct:.0f}%)")
    print(f"{'='*70}")

    # By-specialty breakdown
    spec_results = {}
    for r in results:
        for spec in r.get("expected_specialties", ["Unknown"]):
            if spec not in spec_results:
                spec_results[spec] = {"passed": 0, "total": 0}
            spec_results[spec]["total"] += 1
            if r["passed"]:
                spec_results[spec]["passed"] += 1

    print(f"\n  By specialty:")
    for spec, counts in sorted(spec_results.items()):
        p = counts["passed"]
        t = counts["total"]
        bar = "█" * p + "░" * (t - p)
        print(f"    {spec:25s} {p}/{t} {bar}")

    # Relevance distribution
    all_rels = [r["top_relevance"] for r in results]
    if all_rels:
        avg_rel = sum(all_rels) / len(all_rels)
        min_rel_val = min(all_rels)
        max_rel_val = max(all_rels)
        print(f"\n  Relevance: avg={avg_rel:.3f}  min={min_rel_val:.3f}  max={max_rel_val:.3f}")


if __name__ == "__main__":
    asyncio.run(main())