benchmark.py

#!/usr/bin/env python3
"""
RAG Comparison Benchmark Script
Compares Simple RAG, Agentic RAG, and Graph RAG performance
"""

import json
import time
import argparse
import sys
from pathlib import Path
from typing import Dict, List, Any
import statistics
import os

# Add backend to path
sys.path.insert(0, str(Path(__file__).parent / "backend"))

from app.core.config import Settings
from app.services.embedding_service import EmbeddingService
from app.services.vector_db_service import VectorDBService
from app.services.chunker_service import ChunkerService
from app.services.retrieval_service import RetrievalService
from app.core.constants import RAG_MODES, GROQ_MODELS


class RAGBenchmark:
    """Benchmark harness for RAG modes"""

    def __init__(self, api_key: str = None):
        """Initialize benchmark with Groq API key"""
        self.api_key = api_key or os.getenv("GROQ_API_KEY")
        if not self.api_key:
            raise ValueError("GROQ_API_KEY not set")

        # Initialize services
        self.embedding_service = EmbeddingService("all-MiniLM-L6-v2")
        self.vector_db = VectorDBService("chroma", {"storage_path": "./data/chroma_benchmark"})
        self.chunker = ChunkerService()
        self.retrieval_service = RetrievalService(self.api_key)

        self.results = {
            "timestamp": time.time(),
            "benchmarks": {},
            "summary": {},
        }

    def setup_sample_data(self):
        """Load or create sample test data"""
        sample_docs = [
            {
                "id": "doc_001",
                "title": "Company Overview",
                "content": """
                Our company was founded in 2020 with a mission to revolutionize AI education.
                We provide cutting-edge courses, mentorship, and hands-on projects.

                Core values:
                - Innovation: We push boundaries in AI/ML education
                - Excellence: High-quality, industry-standard curriculum
                - Community: Strong network of learners and mentors
                - Accessibility: Making AI education affordable for everyone

                Our products include online courses, corporate training, and research programs.
                We've trained over 50,000 students across 75 countries.
                """,
            },
            {
                "id": "doc_002",
                "title": "Technical Architecture",
                "content": """
                Our platform uses a microservices architecture with:
                - Frontend: React.js with TypeScript
                - Backend: Python FastAPI
                - Database: PostgreSQL with Redis caching
                - ML: PyTorch and TensorFlow
                - Deployment: Kubernetes on AWS

                Key components:
                1. User Management Service: Handles authentication and profiles
                2. Content Service: Manages courses and materials
                3. Assessment Service: Evaluates student progress
                4. ML Pipeline: Trains and deploys models
                5. Analytics Service: Tracks learning metrics

                All services communicate via REST APIs and message queues.
                """,
            },
            {
                "id": "doc_003",
                "title": "Success Metrics",
                "content": """
                We measure success through:

                Student Outcomes:
                - 95% course completion rate
                - 87% job placement within 6 months
                - Average salary increase: 35%

                Business Metrics:
                - Annual revenue growth: 150%
                - Customer retention: 92%
                - Net Promoter Score: 78

                Learning Metrics:
                - Average improvement: 45% on assessments
                - Skills acquired per student: 8.5
                - Time to competency: 6 months average
                """,
            },
        ]

        print(f"Loading {len(sample_docs)} sample documents...")
        for doc in sample_docs:
            # Create embeddings
            embedding = self.embedding_service.embed(doc["content"])

            # Store in vector DB
            self.vector_db.add_vector(
                vector_id=doc["id"],
                vector=embedding,
                metadata={"title": doc["title"], "doc_id": doc["id"]},
                content=doc["content"],
            )

        print("✓ Sample data loaded")

    def run_benchmark(
        self,
        model_id: str = "llama-3.1-8b-instant",
        rag_modes: List[str] = None,
        iterations: int = 3,
        temperature: float = 0.7,
    ) -> Dict[str, Any]:
        """Run benchmark for specified RAG modes"""

        if rag_modes is None:
            rag_modes = ["simple", "agentic", "graph"]

        test_queries = [
            "What is the company's mission and core values?",
            "Describe the technical architecture and key components.",
            "What are the main success metrics?",
        ]

        results = {
            "model": model_id,
            "rag_modes": {},
            "comparisons": {},
        }

        for mode in rag_modes:
            print(f"\n{'='*60}")
            print(f"Benchmarking: {mode.upper()} RAG")
            print(f"{'='*60}")

            mode_timings = []
            mode_tokens = {"input": [], "output": []}
            mode_sources = []
            mode_costs = []

            for i, query in enumerate(test_queries):
                print(f"  [{i+1}/{len(test_queries)}] Query: {query[:50]}...")

                try:
                    # Get retrieval results
                    search_results = self.vector_db.search(
                        query_embedding=self.embedding_service.embed(query),
                        top_k=5,
                    )

                    # Run RAG mode
                    start_time = time.time()

                    answer, sources, metrics = self.retrieval_service.generate_answer(
                        query=query,
                        search_results=search_results,
                        rag_mode=mode,
                        model_id=model_id,
                        temperature=temperature,
                    )

                    latency = (time.time() - start_time) * 1000

                    # Track metrics
                    mode_timings.append(latency)
                    mode_tokens["input"].append(metrics.get("input_tokens", 0))
                    mode_tokens["output"].append(metrics.get("output_tokens", 0))
                    mode_sources.append(len(sources))
                    mode_costs.append(metrics.get("cost", 0))

                    print(
                        f"    ✓ {latency:.0f}ms | "
                        f"Tokens: {metrics.get('input_tokens', 0)} in, "
                        f"{metrics.get('output_tokens', 0)} out | "
                        f"Sources: {len(sources)}"
                    )

                except Exception as e:
                    print(f"    ✗ Error: {e}")
                    continue

            # Calculate statistics
            if mode_timings:
                results["rag_modes"][mode] = {
                    "latency": {
                        "min_ms": min(mode_timings),
                        "max_ms": max(mode_timings),
                        "mean_ms": statistics.mean(mode_timings),
                        "median_ms": statistics.median(mode_timings),
                        "stdev_ms": statistics.stdev(mode_timings) if len(mode_timings) > 1 else 0,
                    },
                    "tokens": {
                        "input_avg": statistics.mean(mode_tokens["input"]),
                        "output_avg": statistics.mean(mode_tokens["output"]),
                        "total_avg": statistics.mean(
                            [mode_tokens["input"][i] + mode_tokens["output"][i]
                             for i in range(len(mode_tokens["input"]))]
                        ),
                    },
                    "sources_avg": statistics.mean(mode_sources),
                    "cost_per_query_usd": statistics.mean(mode_costs),
                    "cost_1k_queries_usd": statistics.mean(mode_costs) * 1000,
                }

                print(f"\n  Summary for {mode.upper()}:")
                print(
                    f"    Latency: {results['rag_modes'][mode]['latency']['mean_ms']:.0f}ms "
                    f"(±{results['rag_modes'][mode]['latency']['stdev_ms']:.0f}ms)"
                )
                print(
                    f"    Tokens: {results['rag_modes'][mode]['tokens']['input_avg']:.0f} in, "
                    f"{results['rag_modes'][mode]['tokens']['output_avg']:.0f} out"
                )
                print(
                    f"    Sources: {results['rag_modes'][mode]['sources_avg']:.1f} avg"
                )
                print(
                    f"    Cost: ${results['rag_modes'][mode]['cost_per_query_usd']:.4f}/query"
                )

        # Generate comparisons
        results["comparisons"] = self._generate_comparisons(results["rag_modes"])

        return results

    def _generate_comparisons(self, modes_data: Dict[str, Any]) -> Dict[str, Any]:
        """Generate comparative analysis between RAG modes"""
        comparisons = {}

        if len(modes_data) > 1:
            # Latency comparison
            fastest_mode = min(
                modes_data.items(),
                key=lambda x: x[1]["latency"]["mean_ms"],
            )
            comparisons["fastest"] = {
                "mode": fastest_mode[0],
                "latency_ms": fastest_mode[1]["latency"]["mean_ms"],
            }

            # Cost comparison
            cheapest_mode = min(
                modes_data.items(),
                key=lambda x: x[1]["cost_per_query_usd"],
            )
            comparisons["cheapest"] = {
                "mode": cheapest_mode[0],
                "cost_usd": cheapest_mode[1]["cost_per_query_usd"],
            }

            # Most comprehensive (sources)
            most_sources_mode = max(
                modes_data.items(),
                key=lambda x: x[1]["sources_avg"],
            )
            comparisons["most_comprehensive"] = {
                "mode": most_sources_mode[0],
                "sources_avg": most_sources_mode[1]["sources_avg"],
            }

        return comparisons

    def save_results(self, filename: str = "benchmark_results.json"):
        """Save results to JSON file"""
        output_path = Path("data") / filename
        output_path.parent.mkdir(parents=True, exist_ok=True)

        with open(output_path, "w") as f:
            json.dump(self.results, f, indent=2, default=str)

        print(f"\n✓ Results saved to {output_path}")
        return output_path

    def print_summary(self):
        """Print benchmark summary"""
        if not self.results.get("benchmarks"):
            print("No benchmark results yet")
            return

        print(f"\n{'='*60}")
        print("BENCHMARK SUMMARY")
        print(f"{'='*60}\n")

        for model, data in self.results["benchmarks"].items():
            print(f"Model: {model}")
            print(f"  RAG Modes tested: {', '.join(data['rag_modes'].keys())}")

            if data.get("comparisons"):
                print(f"\n  Comparisons:")
                for metric, values in data["comparisons"].items():
                    print(f"    - {metric}: {values['mode']}")

            print()


def main():
    parser = argparse.ArgumentParser(
        description="Benchmark RAG modes (Simple, Agentic, Graph)"
    )
    parser.add_argument(
        "--mode",
        choices=["simple", "agentic", "graph", "all"],
        default="all",
        help="RAG mode(s) to benchmark",
    )
    parser.add_argument(
        "--model",
        default="llama-3.1-8b-instant",
        help="Groq model to use",
    )
    parser.add_argument(
        "--iterations",
        type=int,
        default=3,
        help="Number of test queries per mode",
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.7,
        help="Temperature for generation",
    )
    parser.add_argument(
        "--output",
        default="benchmark_results.json",
        help="Output file for results",
    )

    args = parser.parse_args()

    # Determine modes to benchmark
    if args.mode == "all":
        modes = ["simple", "agentic", "graph"]
    else:
        modes = [args.mode]

    try:
        # Create benchmark
        benchmark = RAGBenchmark()

        # Setup sample data
        benchmark.setup_sample_data()

        # Run benchmark
        print(f"\nStarting benchmark with {args.model}...\n")
        results = benchmark.run_benchmark(
            model_id=args.model,
            rag_modes=modes,
            iterations=args.iterations,
            temperature=args.temperature,
        )

        # Store results
        benchmark.results["benchmarks"][args.model] = results

        # Save and print
        benchmark.save_results(args.output)
        benchmark.print_summary()

        print("\n✅ Benchmark complete!")

    except KeyboardInterrupt:
        print("\n\n⚠ Benchmark interrupted")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()