benchmark_rag_performance.py

#!/usr/bin/env python3
"""
Comprehensive RAG System Performance Benchmark

This script provides detailed performance testing for all components of the RAG system,
including before/after comparisons and detailed metrics analysis.
"""

import asyncio
import gc
import json
import logging
import statistics
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Any

import psutil

# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

try:
    from rag_system.config import RAGConfig
    from rag_system.document_processing.embedder import Embedder, EmbeddingConfig
    from rag_system.generation.generator import Generator
    from rag_system.models import DocumentChunk
    from rag_system.orchestration.langgraph_orchestrator import LangGraphOrchestrator
    from rag_system.retrieval.retriever import Retriever
    from rag_system.storage.qdrant_store import QdrantDocumentStore
except ImportError as e:
    logger.error(f"Failed to import RAG system components: {e}")
    logger.error("Make sure the RAG system is properly installed with: pipx install .")
    sys.exit(1)


class PerformanceMonitor:
    """Monitor system performance during benchmarks."""

    def __init__(self):
        self.start_time = None
        self.start_memory = None
        self.start_cpu = None

    def start(self):
        """Start monitoring."""
        self.start_time = time.time()
        self.start_memory = psutil.virtual_memory().used / 1024 / 1024  # MB
        self.start_cpu = psutil.cpu_percent()

    def stop(self) -> dict[str, float]:
        """Stop monitoring and return metrics."""
        end_time = time.time()
        end_memory = psutil.virtual_memory().used / 1024 / 1024  # MB
        end_cpu = psutil.cpu_percent()

        return {
            "duration": end_time - self.start_time,
            "memory_used_mb": end_memory - self.start_memory,
            "cpu_usage_percent": (self.start_cpu + end_cpu) / 2,
            "peak_memory_mb": psutil.virtual_memory().used / 1024 / 1024,
        }


class RAGBenchmark:
    """Comprehensive RAG system benchmark suite."""

    def __init__(self, config_path: str = "config/rag_config.yaml", save_results: bool = False):
        """Initialize benchmark with configuration."""
        try:
            self.config = RAGConfig.from_file(config_path)
            logger.info(f"Loaded configuration from {config_path}")
        except Exception as e:
            logger.error(f"Failed to load config: {e}")
            logger.info("Using default configuration")
            self.config = RAGConfig()

        self.results = {}
        self.monitor = PerformanceMonitor()
        self.save_results = save_results

        # Test data
        self.test_texts = [
            "Machine learning is a subset of artificial intelligence.",
            "Neural networks are inspired by biological neural networks.",
            "Deep learning uses multiple layers to model data.",
            "Natural language processing enables computers to understand human language.",
            "Computer vision allows machines to interpret visual information.",
            "Reinforcement learning learns through interaction with environment.",
            "Supervised learning uses labeled training data.",
            "Unsupervised learning finds patterns in unlabeled data.",
            "Transfer learning applies knowledge from one domain to another.",
            "Ensemble methods combine multiple models for better performance.",
        ]

        self.test_queries = [
            "What is machine learning?",
            "How do neural networks work?",
            "Explain deep learning concepts",
            "What is natural language processing?",
            "How does computer vision work?",
            "What is reinforcement learning?",
            "Difference between supervised and unsupervised learning",
            "What is transfer learning?",
            "How do ensemble methods work?",
            "Applications of artificial intelligence",
        ]

    async def run_full_benchmark(self) -> dict[str, Any]:
        """Run complete benchmark suite."""
        logger.info("🚀 Starting RAG System Performance Benchmark")
        logger.info("=" * 60)

        start_time = datetime.now()

        # System info
        await self._log_system_info()

        # Component benchmarks
        await self._benchmark_embedding_performance()
        await self._benchmark_storage_performance()
        await self._benchmark_retrieval_performance()
        await self._benchmark_generation_performance()
        await self._benchmark_end_to_end_performance()

        # Cache and optimization tests
        await self._benchmark_cache_performance()
        await self._benchmark_batch_operations()
        await self._benchmark_concurrent_operations()

        # Memory and resource usage
        await self._benchmark_memory_usage()

        end_time = datetime.now()

        # Summary
        self.results["benchmark_info"] = {
            "start_time": start_time.isoformat(),
            "end_time": end_time.isoformat(),
            "total_duration": (end_time - start_time).total_seconds(),
            "config_used": self.config.to_dict(),
        }

        await self._generate_summary()
        return self.results

    async def _log_system_info(self):
        """Log system information."""
        logger.info("📊 System Information")
        logger.info("-" * 30)

        system_info = {
            "cpu_count": psutil.cpu_count(),
            "memory_total_gb": psutil.virtual_memory().total / 1024 / 1024 / 1024,
            "python_version": sys.version,
            "platform": sys.platform,
        }

        for key, value in system_info.items():
            logger.info(f"{key}: {value}")

        self.results["system_info"] = system_info

    async def _benchmark_embedding_performance(self):
        """Benchmark embedding generation performance."""
        logger.info("\n🔤 Benchmarking Embedding Performance")
        logger.info("-" * 40)

        embedding_config = EmbeddingConfig(
            model_name=self.config.embedding_model,
            batch_size=self.config.embedding_batch_size,
            device=self.config.embedding_device,
            normalize_embeddings=self.config.embedding_normalize,
            show_progress_bar=False,
        )

        embedder = Embedder(embedding_config)

        # Single embedding test
        single_times = []
        for i, text in enumerate(self.test_texts[:5]):
            self.monitor.start()
            embedder.embed_text(text)
            metrics = self.monitor.stop()

            single_times.append(metrics["duration"])
            logger.info(f"Single embedding {i + 1}: {metrics['duration']:.3f}s")

        # Batch embedding test
        batch_sizes = [5, 10, 20]
        batch_results = []

        for batch_size in batch_sizes:
            texts = self.test_texts[:batch_size]

            self.monitor.start()
            embeddings = embedder.embed_batch(texts)
            metrics = self.monitor.stop()

            successful = sum(1 for e in embeddings if e is not None)
            throughput = successful / metrics["duration"]

            batch_result = {
                "batch_size": batch_size,
                "duration": metrics["duration"],
                "throughput": throughput,
                "successful_embeddings": successful,
                "memory_used_mb": metrics["memory_used_mb"],
            }
            batch_results.append(batch_result)

            logger.info(
                f"Batch {batch_size}: {metrics['duration']:.3f}s ({throughput:.1f} texts/sec)"
            )

        self.results["embedding_performance"] = {
            "single_embedding": {
                "avg_time": statistics.mean(single_times),
                "min_time": min(single_times),
                "max_time": max(single_times),
                "std_dev": statistics.stdev(single_times) if len(single_times) > 1 else 0,
            },
            "batch_embedding": batch_results,
            "avg_batch_throughput": statistics.mean([b["throughput"] for b in batch_results]),
        }

    async def _benchmark_storage_performance(self):
        """Benchmark vector storage performance."""
        logger.info("\n💾 Benchmarking Storage Performance")
        logger.info("-" * 40)

        document_store = QdrantDocumentStore(self.config)

        # Test connection
        self.monitor.start()
        connected = document_store.connect()
        connection_metrics = self.monitor.stop()

        if not connected:
            logger.warning("Could not connect to Qdrant - skipping storage benchmarks")
            return

        logger.info(f"Connection time: {connection_metrics['duration']:.3f}s")

        # Create test documents
        embedder = Embedder(
            EmbeddingConfig(
                model_name=self.config.embedding_model, batch_size=self.config.embedding_batch_size
            )
        )

        test_docs = []
        embeddings = embedder.embed_batch(self.test_texts)

        for i, (text, embedding) in enumerate(zip(self.test_texts, embeddings)):
            if embedding:
                doc = DocumentChunk(
                    content=text,
                    embedding=embedding,
                    source=f"benchmark_test_{i}.txt",
                    chunk_index=i,
                    metadata={"test": True, "index": i},
                )
                test_docs.append(doc)

        # Test document insertion
        self.monitor.start()
        success, result = document_store.add_documents(test_docs)
        insert_metrics = self.monitor.stop()

        logger.info(
            f"Document insertion: {insert_metrics['duration']:.3f}s ({len(test_docs)} docs)"
        )

        # Test similarity search
        if embeddings[0]:
            search_times = []
            for _i in range(5):
                self.monitor.start()
                document_store.search_similar(
                    query_embedding=embeddings[0], top_k=5, similarity_threshold=0.1
                )
                search_metrics = self.monitor.stop()
                search_times.append(search_metrics["duration"])

            logger.info(f"Avg search time: {statistics.mean(search_times):.3f}s")

        # Test duplicate detection
        duplicate_docs = test_docs[:5] + test_docs[:3]  # Add some duplicates

        self.monitor.start()
        filtered_docs, dup_info = document_store._filter_duplicates(duplicate_docs)
        duplicate_metrics = self.monitor.stop()

        logger.info(f"Duplicate detection: {duplicate_metrics['duration']:.3f}s")
        logger.info(
            f"Found {dup_info['duplicate_count']} duplicates out of {len(duplicate_docs)} docs"
        )

        self.results["storage_performance"] = {
            "connection_time": connection_metrics["duration"],
            "insertion": {
                "duration": insert_metrics["duration"],
                "documents_count": len(test_docs),
                "throughput": len(test_docs) / insert_metrics["duration"],
                "success": success,
            },
            "search": {
                "avg_time": statistics.mean(search_times) if embeddings[0] else 0,
                "min_time": min(search_times) if embeddings[0] else 0,
                "max_time": max(search_times) if embeddings[0] else 0,
            },
            "duplicate_detection": {
                "duration": duplicate_metrics["duration"],
                "input_docs": len(duplicate_docs),
                "duplicates_found": dup_info["duplicate_count"],
                "throughput": len(duplicate_docs) / duplicate_metrics["duration"],
            },
        }

        # Cleanup test documents
        try:
            for _doc in test_docs:
                document_store.client.delete(
                    collection_name=self.config.collection_name,
                    points_selector={
                        "filter": {"must": [{"key": "test", "match": {"value": True}}]}
                    },
                )
        except Exception as e:
            logger.warning(f"Cleanup failed: {e}")

    async def _benchmark_retrieval_performance(self):
        """Benchmark retrieval performance."""
        logger.info("\n🔍 Benchmarking Retrieval Performance")
        logger.info("-" * 40)

        retriever = Retriever(self.config)

        # Single query tests
        single_times = []
        result_counts = []

        for i, query in enumerate(self.test_queries[:5]):
            self.monitor.start()
            try:
                documents = retriever.retrieve(query, top_k=5)
                metrics = self.monitor.stop()

                single_times.append(metrics["duration"])
                result_counts.append(len(documents))

                logger.info(f"Query {i + 1}: {metrics['duration']:.3f}s ({len(documents)} results)")
            except Exception as e:
                logger.warning(f"Query {i + 1} failed: {e}")
                continue

        # Batch query tests
        batch_sizes = [2, 5, 8]
        batch_results = []

        for batch_size in batch_sizes:
            queries = self.test_queries[:batch_size]

            self.monitor.start()
            try:
                results = await retriever.batch_retrieve(queries, top_k=5)
                metrics = self.monitor.stop()

                total_results = sum(len(r) for r in results)
                throughput = len(queries) / metrics["duration"]

                batch_result = {
                    "batch_size": batch_size,
                    "duration": metrics["duration"],
                    "throughput": throughput,
                    "total_results": total_results,
                    "avg_results_per_query": total_results / len(queries),
                }
                batch_results.append(batch_result)

                logger.info(
                    f"Batch {batch_size}: {metrics['duration']:.3f}s ({throughput:.1f} queries/sec)"
                )
            except Exception as e:
                logger.warning(f"Batch retrieval {batch_size} failed: {e}")

        self.results["retrieval_performance"] = {
            "single_query": {
                "avg_time": statistics.mean(single_times) if single_times else 0,
                "min_time": min(single_times) if single_times else 0,
                "max_time": max(single_times) if single_times else 0,
                "avg_results": statistics.mean(result_counts) if result_counts else 0,
            },
            "batch_query": batch_results,
            "avg_batch_throughput": statistics.mean([b["throughput"] for b in batch_results])
            if batch_results
            else 0,
        }

    async def _benchmark_generation_performance(self):
        """Benchmark text generation performance."""
        logger.info("\n🤖 Benchmarking Generation Performance")
        logger.info("-" * 40)

        try:
            generator = Generator(self.config)

            test_contexts = [
                "Machine learning is a subset of artificial intelligence.",
                "Neural networks are computational models inspired by biological neural networks.",
                "Deep learning uses multiple layers to learn representations of data.",
            ]

            generation_times = []

            for i, context in enumerate(test_contexts):
                query = f"Explain this concept: {context[:50]}..."

                self.monitor.start()
                try:
                    response = generator.generate_response(
                        query=query, context_documents=context, max_tokens=100
                    )
                    metrics = self.monitor.stop()

                    generation_times.append(metrics["duration"])
                    response_length = (
                        len(response.answer) if hasattr(response, "answer") else len(str(response))
                    )

                    logger.info(
                        f"Generation {i + 1}: {metrics['duration']:.3f}s ({response_length} chars)"
                    )
                except Exception as e:
                    logger.warning(f"Generation {i + 1} failed: {e}")
                    continue

            self.results["generation_performance"] = {
                "avg_time": statistics.mean(generation_times) if generation_times else 0,
                "min_time": min(generation_times) if generation_times else 0,
                "max_time": max(generation_times) if generation_times else 0,
                "successful_generations": len(generation_times),
            }

        except Exception as e:
            logger.warning(f"Generation benchmark failed: {e}")
            self.results["generation_performance"] = {"error": str(e)}

    async def _benchmark_end_to_end_performance(self):
        """Benchmark complete end-to-end RAG pipeline."""
        logger.info("\n🔄 Benchmarking End-to-End Performance")
        logger.info("-" * 45)

        try:
            orchestrator = LangGraphOrchestrator(self.config)

            e2e_times = []

            for i, query in enumerate(self.test_queries[:3]):
                self.monitor.start()
                try:
                    await orchestrator.process_query_async(query)
                    metrics = self.monitor.stop()

                    e2e_times.append(metrics["duration"])

                    logger.info(f"E2E Query {i + 1}: {metrics['duration']:.3f}s")
                except Exception as e:
                    logger.warning(f"E2E Query {i + 1} failed: {e}")
                    continue

            self.results["end_to_end_performance"] = {
                "avg_time": statistics.mean(e2e_times) if e2e_times else 0,
                "min_time": min(e2e_times) if e2e_times else 0,
                "max_time": max(e2e_times) if e2e_times else 0,
                "successful_queries": len(e2e_times),
            }

        except Exception as e:
            logger.warning(f"End-to-end benchmark failed: {e}")
            self.results["end_to_end_performance"] = {"error": str(e)}

    async def _benchmark_cache_performance(self):
        """Benchmark caching effectiveness."""
        logger.info("\n⚡ Benchmarking Cache Performance")
        logger.info("-" * 40)

        retriever = Retriever(self.config)
        test_query = "What is machine learning?"

        # First query (cache miss)
        self.monitor.start()
        try:
            results1 = retriever.retrieve(test_query)
            miss_metrics = self.monitor.stop()

            # Second identical query (cache hit)
            self.monitor.start()
            results2 = retriever.retrieve(test_query)
            hit_metrics = self.monitor.stop()

            speedup = (
                miss_metrics["duration"] / hit_metrics["duration"]
                if hit_metrics["duration"] > 0
                else float("inf")
            )

            logger.info(f"Cache miss: {miss_metrics['duration']:.3f}s")
            logger.info(f"Cache hit: {hit_metrics['duration']:.3f}s")
            logger.info(f"Speedup: {speedup:.1f}x")

            self.results["cache_performance"] = {
                "cache_miss_time": miss_metrics["duration"],
                "cache_hit_time": hit_metrics["duration"],
                "speedup_factor": speedup,
                "results_identical": len(results1) == len(results2),
                "cache_enabled": self.config.retrieval_enable_cache,
            }

        except Exception as e:
            logger.warning(f"Cache benchmark failed: {e}")
            self.results["cache_performance"] = {"error": str(e)}

    async def _benchmark_batch_operations(self):
        """Benchmark batch vs individual operations."""
        logger.info("\n📦 Benchmarking Batch Operations")
        logger.info("-" * 40)

        embedder = Embedder(
            EmbeddingConfig(
                model_name=self.config.embedding_model, batch_size=self.config.embedding_batch_size
            )
        )

        test_texts = self.test_texts[:10]

        # Individual embeddings
        self.monitor.start()
        individual_embeddings = []
        for text in test_texts:
            embedding = embedder.embed_text(text)
            individual_embeddings.append(embedding)
        individual_metrics = self.monitor.stop()

        # Batch embeddings
        self.monitor.start()
        embedder.embed_batch(test_texts)
        batch_metrics = self.monitor.stop()

        speedup = individual_metrics["duration"] / batch_metrics["duration"]

        logger.info(f"Individual: {individual_metrics['duration']:.3f}s")
        logger.info(f"Batch: {batch_metrics['duration']:.3f}s")
        logger.info(f"Batch speedup: {speedup:.1f}x")

        self.results["batch_operations"] = {
            "individual_time": individual_metrics["duration"],
            "batch_time": batch_metrics["duration"],
            "speedup_factor": speedup,
            "batch_size": len(test_texts),
            "individual_memory_mb": individual_metrics["memory_used_mb"],
            "batch_memory_mb": batch_metrics["memory_used_mb"],
        }

    async def _benchmark_concurrent_operations(self):
        """Benchmark concurrent operations."""
        logger.info("\n🔀 Benchmarking Concurrent Operations")
        logger.info("-" * 45)

        retriever = Retriever(self.config)
        queries = self.test_queries[:5]

        # Sequential processing
        self.monitor.start()
        sequential_results = []
        for query in queries:
            try:
                result = retriever.retrieve(query, top_k=3)
                sequential_results.append(result)
            except Exception as e:
                logger.warning(f"Sequential query failed: {e}")
        sequential_metrics = self.monitor.stop()

        # Concurrent processing
        self.monitor.start()
        try:
            await retriever.batch_retrieve(queries, top_k=3)
            concurrent_metrics = self.monitor.stop()

            speedup = sequential_metrics["duration"] / concurrent_metrics["duration"]

            logger.info(f"Sequential: {sequential_metrics['duration']:.3f}s")
            logger.info(f"Concurrent: {concurrent_metrics['duration']:.3f}s")
            logger.info(f"Concurrent speedup: {speedup:.1f}x")

            self.results["concurrent_operations"] = {
                "sequential_time": sequential_metrics["duration"],
                "concurrent_time": concurrent_metrics["duration"],
                "speedup_factor": speedup,
                "query_count": len(queries),
                "sequential_memory_mb": sequential_metrics["memory_used_mb"],
                "concurrent_memory_mb": concurrent_metrics["memory_used_mb"],
            }
        except Exception as e:
            logger.warning(f"Concurrent benchmark failed: {e}")
            self.results["concurrent_operations"] = {"error": str(e)}

    async def _benchmark_memory_usage(self):
        """Benchmark memory usage patterns."""
        logger.info("\n🧠 Benchmarking Memory Usage")
        logger.info("-" * 35)

        # Force garbage collection
        gc.collect()
        initial_memory = psutil.virtual_memory().used / 1024 / 1024

        embedder = Embedder(
            EmbeddingConfig(
                model_name=self.config.embedding_model, batch_size=self.config.embedding_batch_size
            )
        )

        # Memory usage during embedding
        psutil.virtual_memory().used / 1024 / 1024
        large_texts = self.test_texts * 10  # 100 texts
        embeddings = embedder.embed_batch(large_texts)
        memory_after = psutil.virtual_memory().used / 1024 / 1024

        # Force cleanup
        del embeddings
        gc.collect()
        memory_cleaned = psutil.virtual_memory().used / 1024 / 1024

        logger.info(f"Initial memory: {initial_memory:.1f} MB")
        logger.info(f"After embedding: {memory_after:.1f} MB")
        logger.info(f"After cleanup: {memory_cleaned:.1f} MB")
        logger.info(f"Peak usage: {memory_after - initial_memory:.1f} MB")

        self.results["memory_usage"] = {
            "initial_memory_mb": initial_memory,
            "peak_memory_mb": memory_after,
            "after_cleanup_mb": memory_cleaned,
            "peak_usage_mb": memory_after - initial_memory,
            "cleanup_efficiency": (memory_after - memory_cleaned)
            / (memory_after - initial_memory)
            * 100,
        }

    async def _generate_summary(self):
        """Generate benchmark summary."""
        logger.info("\n" + "=" * 60)
        logger.info("📊 BENCHMARK SUMMARY")
        logger.info("=" * 60)

        # Performance highlights
        highlights = []

        if "embedding_performance" in self.results:
            ep = self.results["embedding_performance"]
            if "avg_batch_throughput" in ep:
                highlights.append(
                    f"Embedding throughput: {ep['avg_batch_throughput']:.1f} texts/sec"
                )

        if (
            "cache_performance" in self.results
            and "speedup_factor" in self.results["cache_performance"]
        ):
            cp = self.results["cache_performance"]
            highlights.append(f"Cache speedup: {cp['speedup_factor']:.1f}x")

        if (
            "batch_operations" in self.results
            and "speedup_factor" in self.results["batch_operations"]
        ):
            bo = self.results["batch_operations"]
            highlights.append(f"Batch speedup: {bo['speedup_factor']:.1f}x")

        if (
            "concurrent_operations" in self.results
            and "speedup_factor" in self.results["concurrent_operations"]
        ):
            co = self.results["concurrent_operations"]
            highlights.append(f"Concurrent speedup: {co['speedup_factor']:.1f}x")

        logger.info("🎯 Performance Highlights:")
        for highlight in highlights:
            logger.info(f"  • {highlight}")

        # Configuration summary
        logger.info("\n⚙️  Configuration Used:")
        logger.info(f"  • Embedding batch size: {self.config.embedding_batch_size}")
        logger.info(f"  • Cache enabled: {self.config.retrieval_enable_cache}")
        logger.info(f"  • gRPC enabled: {self.config.qdrant_prefer_grpc}")
        logger.info(f"  • Max concurrent searches: {self.config.retrieval_max_concurrent_searches}")

        # Save results only if requested
        if self.save_results:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            results_file = f"benchmark_results_{timestamp}.json"

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

            logger.info(f"\n💾 Results saved to: {results_file}")
        else:
            logger.info("\n💾 Results not saved (use --save-results to save)")

        logger.info("=" * 60)


async def main():
    """Run the RAG system benchmark."""
    import argparse

    parser = argparse.ArgumentParser(description="RAG System Performance Benchmark")
    parser.add_argument(
        "--config",
        "-c",
        type=str,
        default="config/rag_config.yaml",
        help="Path to configuration file",
    )
    parser.add_argument(
        "--save-results", "-s", action="store_true", help="Save benchmark results to JSON file"
    )

    args = parser.parse_args()

    print("🚀 RAG System Performance Benchmark")
    print("=" * 50)

    # Check if config file exists
    config_path = args.config
    if not Path(config_path).exists():
        print(f"⚠️  Configuration file not found: {config_path}")
        print("Using default configuration...")
        config_path = None

    try:
        benchmark = (
            RAGBenchmark(config_path, save_results=args.save_results)
            if config_path
            else RAGBenchmark(save_results=args.save_results)
        )
        results = await benchmark.run_full_benchmark()

        print("\n✅ Benchmark completed successfully!")
        print(f"📊 Total tests run: {len([k for k in results if k.endswith('_performance')])}")

        if args.save_results:
            print("💾 Results saved to JSON file")
        else:
            print("💡 Use --save-results flag to save detailed results to JSON file")

        return results

    except KeyboardInterrupt:
        print("\n⚠️  Benchmark interrupted by user")
        return None
    except Exception as e:
        print(f"\n❌ Benchmark failed: {e}")
        logger.exception("Benchmark error details:")
        return None


if __name__ == "__main__":
    # Run the benchmark
    results = asyncio.run(main())

    if results:
        print("\n🎉 Benchmark completed! Check the results file for detailed metrics.")
    else:
        print("\n💥 Benchmark failed or was interrupted.")
        sys.exit(1)