evaluate.py

"""
Evaluation Framework
====================
Runs a suite of test queries against the multi-agent pipeline and measures:
  - Retrieval accuracy (does the correct source appear in top-k?)
  - Answer relevance (keyword overlap with expected answer)
  - Latency per stage (retrieval, summarization, response)
  - End-to-end latency

Generates a JSON report and prints a summary table.
"""

import json
import os
import time
from dataclasses import dataclass, asdict
from typing import Optional

from agents import MultiAgentCoordinator, AgentConfig


@dataclass
class EvalCase:
    """A single test case."""
    query: str
    expected_source: str  # which document should be retrieved
    expected_keywords: list[str]  # keywords that should appear in the answer


@dataclass
class EvalResult:
    """Result of running one eval case."""
    query: str
    source_hit: bool  # did expected source appear in retrieved chunks?
    source_rank: int | None  # rank of expected source (0-indexed), None if missed
    keyword_recall: float  # fraction of expected keywords found in answer
    matched_keywords: list[str]
    missed_keywords: list[str]
    retrieval_ms: float
    summarization_ms: float
    response_ms: float
    total_ms: float
    answer_preview: str  # first 200 chars


def build_eval_cases() -> list[EvalCase]:
    """
    Define test queries mapped to the sample documents.
    These align with the sample docs in data/.
    """
    return [
        # --- machine_learning_basics.txt ---
        EvalCase(
            query="What is supervised learning?",
            expected_source="machine_learning_basics.txt",
            expected_keywords=["labeled", "training", "prediction"],
        ),
        EvalCase(
            query="Explain the difference between classification and regression.",
            expected_source="machine_learning_basics.txt",
            expected_keywords=["classification", "regression", "categories", "continuous"],
        ),
        EvalCase(
            query="What is overfitting and how can it be prevented?",
            expected_source="machine_learning_basics.txt",
            expected_keywords=["overfitting", "regularization", "training"],
        ),
        EvalCase(
            query="What are common evaluation metrics for classification models?",
            expected_source="machine_learning_basics.txt",
            expected_keywords=["accuracy", "precision", "recall"],
        ),
        EvalCase(
            query="How does gradient descent work?",
            expected_source="machine_learning_basics.txt",
            expected_keywords=["gradient", "loss", "parameters", "minimize"],
        ),

        # --- transformer_architecture.txt ---
        EvalCase(
            query="What is the attention mechanism in transformers?",
            expected_source="transformer_architecture.txt",
            expected_keywords=["attention", "query", "key", "value"],
        ),
        EvalCase(
            query="How does self-attention differ from cross-attention?",
            expected_source="transformer_architecture.txt",
            expected_keywords=["self-attention", "cross-attention", "sequence"],
        ),
        EvalCase(
            query="What is positional encoding and why is it needed?",
            expected_source="transformer_architecture.txt",
            expected_keywords=["positional", "encoding", "order", "position"],
        ),
        EvalCase(
            query="Explain the encoder-decoder structure of transformers.",
            expected_source="transformer_architecture.txt",
            expected_keywords=["encoder", "decoder", "layers"],
        ),
        EvalCase(
            query="What are multi-head attention mechanisms?",
            expected_source="transformer_architecture.txt",
            expected_keywords=["multi-head", "attention", "heads", "parallel"],
        ),

        # --- rag_systems.txt ---
        EvalCase(
            query="What is Retrieval Augmented Generation?",
            expected_source="rag_systems.txt",
            expected_keywords=["retrieval", "generation", "knowledge", "documents"],
        ),
        EvalCase(
            query="How do vector embeddings work in RAG systems?",
            expected_source="rag_systems.txt",
            expected_keywords=["vector", "embedding", "similarity", "search"],
        ),
        EvalCase(
            query="What are the main challenges of RAG pipelines?",
            expected_source="rag_systems.txt",
            expected_keywords=["challenges", "retrieval", "quality"],
        ),
        EvalCase(
            query="How does chunking strategy affect RAG performance?",
            expected_source="rag_systems.txt",
            expected_keywords=["chunk", "size", "overlap"],
        ),
        EvalCase(
            query="What vector databases are commonly used for RAG?",
            expected_source="rag_systems.txt",
            expected_keywords=["vector", "database", "FAISS"],
        ),

        # --- llm_agents.txt ---
        EvalCase(
            query="What is an LLM agent and how does it differ from a chatbot?",
            expected_source="llm_agents.txt",
            expected_keywords=["agent", "tools", "actions", "reasoning"],
        ),
        EvalCase(
            query="What is the ReAct pattern for LLM agents?",
            expected_source="llm_agents.txt",
            expected_keywords=["ReAct", "reasoning", "action"],
        ),
        EvalCase(
            query="How do multi-agent systems coordinate?",
            expected_source="llm_agents.txt",
            expected_keywords=["multi-agent", "coordination", "communication"],
        ),
        EvalCase(
            query="What are tool-use capabilities in LLM agents?",
            expected_source="llm_agents.txt",
            expected_keywords=["tool", "function", "calling", "API"],
        ),
        EvalCase(
            query="What is agent memory and why does it matter?",
            expected_source="llm_agents.txt",
            expected_keywords=["memory", "context", "history", "state"],
        ),

        # --- prompt_engineering.txt ---
        EvalCase(
            query="What are effective prompt engineering techniques?",
            expected_source="prompt_engineering.txt",
            expected_keywords=["prompt", "few-shot", "chain-of-thought"],
        ),
        EvalCase(
            query="How does chain-of-thought prompting improve reasoning?",
            expected_source="prompt_engineering.txt",
            expected_keywords=["chain-of-thought", "reasoning", "step"],
        ),
        EvalCase(
            query="What is few-shot prompting?",
            expected_source="prompt_engineering.txt",
            expected_keywords=["few-shot", "examples", "prompt"],
        ),
        EvalCase(
            query="How can prompts be optimized for structured output?",
            expected_source="prompt_engineering.txt",
            expected_keywords=["structured", "output", "format", "JSON"],
        ),
        EvalCase(
            query="What is the role of system prompts?",
            expected_source="prompt_engineering.txt",
            expected_keywords=["system", "prompt", "behavior", "instruction"],
        ),

        # --- Cross-document queries ---
        EvalCase(
            query="How do transformers relate to RAG systems?",
            expected_source="rag_systems.txt",  # primary source
            expected_keywords=["transformer", "retrieval", "generation", "embedding"],
        ),
        EvalCase(
            query="What role does prompt engineering play in agent workflows?",
            expected_source="llm_agents.txt",
            expected_keywords=["prompt", "agent", "instruction"],
        ),
    ]


def run_eval_case(coordinator: MultiAgentCoordinator, case: EvalCase) -> EvalResult:
    """Run a single evaluation case through the pipeline with timing."""

    # Timed retrieval
    t0 = time.time()
    chunks = coordinator._retrieval_agent.retrieve(case.query)
    retrieval_ms = (time.time() - t0) * 1000

    # Check source hit
    retrieved_sources = [c["source"] for c in chunks]
    source_hit = case.expected_source in retrieved_sources
    source_rank = retrieved_sources.index(case.expected_source) if source_hit else None

    # Timed summarization
    t0 = time.time()
    summary = coordinator._summarization_agent.summarize(case.query, chunks)
    summarization_ms = (time.time() - t0) * 1000

    # Timed response
    sources = list({c["source"] for c in chunks})
    t0 = time.time()
    answer = coordinator._response_agent.generate(case.query, summary, sources)
    response_ms = (time.time() - t0) * 1000

    # Keyword recall
    answer_lower = answer.lower()
    matched = [kw for kw in case.expected_keywords if kw.lower() in answer_lower]
    missed = [kw for kw in case.expected_keywords if kw.lower() not in answer_lower]
    recall = len(matched) / len(case.expected_keywords) if case.expected_keywords else 1.0

    return EvalResult(
        query=case.query,
        source_hit=source_hit,
        source_rank=source_rank,
        keyword_recall=round(recall, 3),
        matched_keywords=matched,
        missed_keywords=missed,
        retrieval_ms=round(retrieval_ms, 1),
        summarization_ms=round(summarization_ms, 1),
        response_ms=round(response_ms, 1),
        total_ms=round(retrieval_ms + summarization_ms + response_ms, 1),
        answer_preview=answer[:200],
    )


def run_evaluation(data_dir: str = "data", report_path: str = "eval_report.json"):
    """Run the full evaluation suite."""
    config = AgentConfig()
    coordinator = MultiAgentCoordinator(config)

    # Ingest sample docs
    import glob
    doc_files = sorted(glob.glob(os.path.join(data_dir, "*.txt")))
    if not doc_files:
        print(f"ERROR: No .txt files found in {data_dir}/")
        print("Run: python create_sample_docs.py")
        return

    print(f"Ingesting {len(doc_files)} documents...")
    coordinator.ingest(file_paths=doc_files)

    # Run eval cases
    cases = build_eval_cases()
    print(f"\nRunning {len(cases)} evaluation queries...\n")

    results: list[EvalResult] = []
    for i, case in enumerate(cases, 1):
        print(f"[{i:2d}/{len(cases)}] {case.query[:60]}...", end=" ", flush=True)
        result = run_eval_case(coordinator, case)
        status = "✓" if result.source_hit else "✗"
        print(f"{status} (recall={result.keyword_recall:.0%}, {result.total_ms:.0f}ms)")
        results.append(result)

    # Summary
    print("\n" + "=" * 60)
    print("EVALUATION SUMMARY")
    print("=" * 60)

    total = len(results)
    source_hits = sum(1 for r in results if r.source_hit)
    avg_recall = sum(r.keyword_recall for r in results) / total
    avg_latency = sum(r.total_ms for r in results) / total
    avg_retrieval = sum(r.retrieval_ms for r in results) / total
    avg_summarization = sum(r.summarization_ms for r in results) / total
    avg_response = sum(r.response_ms for r in results) / total

    print(f"  Source Retrieval Accuracy:  {source_hits}/{total} ({source_hits/total:.0%})")
    print(f"  Average Keyword Recall:     {avg_recall:.1%}")
    print(f"  Average Latency (total):    {avg_latency:.0f}ms")
    print(f"    - Retrieval:              {avg_retrieval:.0f}ms")
    print(f"    - Summarization:          {avg_summarization:.0f}ms")
    print(f"    - Response Generation:    {avg_response:.0f}ms")

    # Failures
    failures = [r for r in results if not r.source_hit]
    if failures:
        print(f"\n  Failed retrievals ({len(failures)}):")
        for f in failures:
            print(f"    - {f.query[:60]}")

    # Save report
    report = {
        "summary": {
            "total_queries": total,
            "source_accuracy": round(source_hits / total, 3),
            "avg_keyword_recall": round(avg_recall, 3),
            "avg_total_ms": round(avg_latency, 1),
            "avg_retrieval_ms": round(avg_retrieval, 1),
            "avg_summarization_ms": round(avg_summarization, 1),
            "avg_response_ms": round(avg_response, 1),
        },
        "results": [asdict(r) for r in results],
    }
    with open(report_path, "w") as f:
        json.dump(report, f, indent=2)
    print(f"\nFull report saved to {report_path}")


if __name__ == "__main__":
    run_evaluation()