Embedding Optimization

"""
Large-Scale Document Batch Processing

Process thousands to millions of documents efficiently with parallel processing,
rate limiting, and error recovery.

Dependencies:
    pip install openai numpy tqdm

Usage:
    python batch_processor.py --input documents.jsonl --output embeddings.npy

Environment Variables:
    OPENAI_API_KEY: Your OpenAI API key
"""

import os
import json
import time
import argparse
import numpy as np
from typing import List, Dict, Optional, Iterator
from dataclasses import dataclass
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
import hashlib


@dataclass
class ProcessingStats:
    """Statistics for batch processing."""

    total_documents: int = 0
    successful: int = 0
    failed: int = 0
    skipped: int = 0
    total_tokens: int = 0
    total_cost: float = 0.0
    start_time: float = 0.0
    end_time: float = 0.0

    @property
    def elapsed_seconds(self) -> float:
        """Get elapsed time in seconds."""
        return self.end_time - self.start_time if self.end_time > 0 else 0

    @property
    def throughput(self) -> float:
        """Get processing throughput (docs/sec)."""
        if self.elapsed_seconds > 0:
            return self.successful / self.elapsed_seconds
        return 0.0


class RateLimiter:
    """Rate limiter for API calls."""

    def __init__(self, max_requests_per_minute: int = 500):
        """
        Initialize rate limiter.

        Args:
            max_requests_per_minute: Maximum API requests per minute
        """
        self.max_requests = max_requests_per_minute
        self.window_seconds = 60
        self.requests = []

    def acquire(self):
        """Wait if necessary to respect rate limit."""
        now = time.time()

        # Remove requests outside current window
        self.requests = [t for t in self.requests if now - t < self.window_seconds]

        # Wait if at limit
        if len(self.requests) >= self.max_requests:
            oldest = min(self.requests)
            sleep_time = self.window_seconds - (now - oldest)
            if sleep_time > 0:
                time.sleep(sleep_time)

        # Record this request
        self.requests.append(time.time())


class BatchProcessor:
    """Process large document collections with embeddings."""

    def __init__(
        self,
        model: str = 'text-embedding-3-small',
        batch_size: int = 100,
        max_workers: int = 5,
        max_requests_per_minute: int = 500,
        retry_attempts: int = 3,
        retry_delay: float = 1.0
    ):
        """
        Initialize batch processor.

        Args:
            model: OpenAI embedding model
            batch_size: Number of documents per API call
            max_workers: Number of parallel workers
            max_requests_per_minute: API rate limit
            retry_attempts: Number of retries on failure
            retry_delay: Delay between retries (seconds)
        """
        try:
            from openai import OpenAI
            self.client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
        except ImportError:
            raise ImportError("Install openai: pip install openai")

        self.model = model
        self.batch_size = batch_size
        self.max_workers = max_workers
        self.retry_attempts = retry_attempts
        self.retry_delay = retry_delay

        self.rate_limiter = RateLimiter(max_requests_per_minute)
        self.stats = ProcessingStats()

        # Cost per 1M tokens
        self.cost_per_1m = {
            'text-embedding-3-small': 0.02,
            'text-embedding-3-large': 0.13,
        }.get(model, 0.02)

    def _estimate_tokens(self, text: str) -> int:
        """Estimate token count (rough approximation)."""
        return len(text) // 4

    def _embed_batch(self, texts: List[str]) -> List[List[float]]:
        """
        Embed a batch of texts with retry logic.

        Args:
            texts: List of texts to embed

        Returns:
            List of embedding vectors
        """
        for attempt in range(self.retry_attempts):
            try:
                # Respect rate limit
                self.rate_limiter.acquire()

                # Call API
                response = self.client.embeddings.create(
                    model=self.model,
                    input=texts
                )

                # Extract embeddings
                embeddings = [d.embedding for d in response.data]

                # Update stats
                for text in texts:
                    tokens = self._estimate_tokens(text)
                    self.stats.total_tokens += tokens
                    self.stats.total_cost += (tokens / 1_000_000) * self.cost_per_1m

                return embeddings

            except Exception as e:
                if attempt < self.retry_attempts - 1:
                    # Exponential backoff
                    delay = self.retry_delay * (2 ** attempt)
                    print(f"  Retry {attempt + 1}/{self.retry_attempts} after {delay:.1f}s: {e}")
                    time.sleep(delay)
                else:
                    print(f"  Failed after {self.retry_attempts} attempts: {e}")
                    raise

    def _process_chunk(self, documents: List[Dict]) -> List[Optional[List[float]]]:
        """
        Process a chunk of documents.

        Args:
            documents: List of document dicts with 'text' field

        Returns:
            List of embeddings (None for failed documents)
        """
        texts = [doc.get('text', '') for doc in documents]

        try:
            embeddings = self._embed_batch(texts)
            self.stats.successful += len(documents)
            return embeddings
        except Exception as e:
            self.stats.failed += len(documents)
            return [None] * len(documents)

    def load_documents(self, input_path: Path) -> Iterator[List[Dict]]:
        """
        Load documents from JSONL file in batches.

        Args:
            input_path: Path to JSONL file

        Yields:
            Batches of documents
        """
        batch = []

        with open(input_path, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip():
                    doc = json.loads(line)
                    batch.append(doc)

                    if len(batch) >= self.batch_size:
                        yield batch
                        batch = []

        # Yield remaining documents
        if batch:
            yield batch

    def process_documents(
        self,
        input_path: Path,
        output_path: Path,
        resume: bool = False
    ) -> ProcessingStats:
        """
        Process documents from JSONL file and save embeddings.

        Args:
            input_path: Path to input JSONL file
            output_path: Path to output .npy file
            resume: Resume from checkpoint if available

        Returns:
            Processing statistics
        """
        print(f"Processing documents from {input_path}")
        print(f"Model: {self.model}")
        print(f"Batch size: {self.batch_size}")
        print(f"Max workers: {self.max_workers}")
        print(f"Rate limit: {self.rate_limiter.max_requests} req/min")
        print()

        # Count total documents
        with open(input_path, 'r') as f:
            total_docs = sum(1 for line in f if line.strip())
        self.stats.total_documents = total_docs
        print(f"Total documents: {total_docs:,}")

        # Check for checkpoint
        checkpoint_path = output_path.with_suffix('.checkpoint.npy')
        processed_docs = 0
        all_embeddings = []

        if resume and checkpoint_path.exists():
            print(f"\nResuming from checkpoint: {checkpoint_path}")
            all_embeddings = np.load(checkpoint_path).tolist()
            processed_docs = len(all_embeddings)
            print(f"Loaded {processed_docs:,} previously processed embeddings")

        # Initialize stats
        self.stats.start_time = time.time()

        try:
            # Progress tracking
            try:
                from tqdm import tqdm
                progress = tqdm(total=total_docs, initial=processed_docs, desc="Processing")
            except ImportError:
                progress = None
                print("Install tqdm for progress bar: pip install tqdm")

            # Process in parallel batches
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                futures = []

                # Submit batches
                for batch_docs in self.load_documents(input_path):
                    # Skip if already processed (for resume)
                    if processed_docs >= len(all_embeddings) + len(batch_docs):
                        processed_docs += len(batch_docs)
                        if progress:
                            progress.update(len(batch_docs))
                        continue

                    future = executor.submit(self._process_chunk, batch_docs)
                    futures.append(future)

                    # Limit queue size to avoid memory issues
                    if len(futures) >= self.max_workers * 2:
                        # Wait for some to complete
                        done = as_completed(futures[:self.max_workers])
                        for future in done:
                            embeddings = future.result()
                            all_embeddings.extend(embeddings)
                            if progress:
                                progress.update(len(embeddings))

                        futures = futures[self.max_workers:]

                # Wait for remaining futures
                for future in as_completed(futures):
                    embeddings = future.result()
                    all_embeddings.extend(embeddings)
                    if progress:
                        progress.update(len(embeddings))

            if progress:
                progress.close()

            # Save results
            self.stats.end_time = time.time()

            # Filter out failed embeddings (None values)
            valid_embeddings = [e for e in all_embeddings if e is not None]

            if valid_embeddings:
                embeddings_array = np.array(valid_embeddings, dtype=np.float32)
                np.save(output_path, embeddings_array)
                print(f"\nSaved {len(valid_embeddings):,} embeddings to {output_path}")

                # Remove checkpoint if exists
                if checkpoint_path.exists():
                    checkpoint_path.unlink()
            else:
                print("\nNo valid embeddings generated")

        except KeyboardInterrupt:
            print("\n\nInterrupted! Saving checkpoint...")
            if all_embeddings:
                checkpoint_array = np.array([e for e in all_embeddings if e is not None], dtype=np.float32)
                np.save(checkpoint_path, checkpoint_array)
                print(f"Checkpoint saved to {checkpoint_path}")
                print("Resume with --resume flag")
            raise

        return self.stats

    def print_summary(self):
        """Print processing summary."""
        print("\n" + "=" * 60)
        print("BATCH PROCESSING SUMMARY")
        print("=" * 60)

        print(f"\nDocuments:")
        print(f"  Total:       {self.stats.total_documents:,}")
        print(f"  Successful:  {self.stats.successful:,}")
        print(f"  Failed:      {self.stats.failed:,}")
        print(f"  Skipped:     {self.stats.skipped:,}")

        success_rate = (self.stats.successful / self.stats.total_documents * 100) if self.stats.total_documents > 0 else 0
        print(f"  Success Rate: {success_rate:.2f}%")

        print(f"\nPerformance:")
        print(f"  Elapsed Time: {self.stats.elapsed_seconds:.2f}s")
        print(f"  Throughput:   {self.stats.throughput:.2f} docs/sec")

        print(f"\nCost:")
        print(f"  Total Tokens: {self.stats.total_tokens:,}")
        print(f"  Total Cost:   ${self.stats.total_cost:.4f}")
        print(f"  Cost per Doc: ${self.stats.total_cost / max(1, self.stats.successful):.6f}")

        print("=" * 60)


def main():
    """CLI for batch processing."""

    parser = argparse.ArgumentParser(
        description='Batch process documents with embeddings',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument(
        '--input',
        type=Path,
        required=True,
        help='Input JSONL file (one document per line with "text" field)'
    )

    parser.add_argument(
        '--output',
        type=Path,
        required=True,
        help='Output .npy file for embeddings'
    )

    parser.add_argument(
        '--model',
        type=str,
        default='text-embedding-3-small',
        choices=['text-embedding-3-small', 'text-embedding-3-large'],
        help='OpenAI embedding model'
    )

    parser.add_argument(
        '--batch-size',
        type=int,
        default=100,
        help='Documents per API call (max 2048)'
    )

    parser.add_argument(
        '--workers',
        type=int,
        default=5,
        help='Number of parallel workers'
    )

    parser.add_argument(
        '--rate-limit',
        type=int,
        default=500,
        help='Max API requests per minute'
    )

    parser.add_argument(
        '--resume',
        action='store_true',
        help='Resume from checkpoint if available'
    )

    args = parser.parse_args()

    # Validate inputs
    if not args.input.exists():
        print(f"Error: Input file not found: {args.input}")
        return

    if args.batch_size > 2048:
        print("Warning: OpenAI max batch size is 2048")
        args.batch_size = 2048

    # Create processor
    processor = BatchProcessor(
        model=args.model,
        batch_size=args.batch_size,
        max_workers=args.workers,
        max_requests_per_minute=args.rate_limit
    )

    # Process documents
    try:
        stats = processor.process_documents(
            input_path=args.input,
            output_path=args.output,
            resume=args.resume
        )

        processor.print_summary()

        print("\nProcessing complete!")
        print(f"\nLoad embeddings with:")
        print(f"  import numpy as np")
        print(f"  embeddings = np.load('{args.output}')")

    except KeyboardInterrupt:
        print("\n\nProcessing interrupted by user")
    except Exception as e:
        print(f"\nError: {e}")
        raise


if __name__ == "__main__":
    main()

"""
Embedding Model Benchmarking

Benchmark different embedding models for latency, quality, and cost.
Compare OpenAI, Cohere, Voyage AI, and open-source models.

Dependencies:
    pip install numpy sentence-transformers openai cohere voyageai matplotlib

Usage:
    python benchmark_embeddings.py

Environment Variables:
    OPENAI_API_KEY: Your OpenAI API key (optional)
    COHERE_API_KEY: Your Cohere API key (optional)
    VOYAGE_API_KEY: Your Voyage AI API key (optional)
"""

import os
import time
import numpy as np
from typing import List, Dict, Optional, Callable
from dataclasses import dataclass, field
import matplotlib.pyplot as plt
from datetime import datetime


@dataclass
class BenchmarkResult:
    """Results from benchmarking a single model."""

    model_name: str
    provider: str  # 'openai', 'cohere', 'voyage', 'local'
    dimensions: int

    # Latency metrics (milliseconds)
    latency_p50: float = 0.0
    latency_p95: float = 0.0
    latency_p99: float = 0.0
    latency_avg: float = 0.0

    # Throughput
    throughput_texts_per_sec: float = 0.0

    # Cost (per 1M tokens)
    cost_per_1m_tokens: float = 0.0

    # Quality (MTEB score approximation or known benchmark)
    mteb_score: Optional[float] = None

    # Additional metadata
    batch_size: int = 1
    num_samples: int = 0
    errors: int = 0

    def cost_per_1k_texts(self, avg_tokens_per_text: int = 100) -> float:
        """Calculate cost per 1,000 texts."""
        if self.cost_per_1m_tokens == 0:
            return 0.0
        return (avg_tokens_per_text * 1000 / 1_000_000) * self.cost_per_1m_tokens

    def quality_per_dollar(self, avg_tokens_per_text: int = 100) -> float:
        """Calculate quality score per dollar (higher is better)."""
        if not self.mteb_score or self.cost_per_1m_tokens == 0:
            return 0.0
        cost_per_1k = self.cost_per_1k_texts(avg_tokens_per_text)
        if cost_per_1k == 0:
            return float('inf')  # Local models
        return self.mteb_score / cost_per_1k


class EmbedderBenchmark:
    """Benchmark embedding models."""

    # Known MTEB scores (as of Nov 2024)
    MTEB_SCORES = {
        'text-embedding-3-small': 62.3,
        'text-embedding-3-large': 64.6,
        'embed-english-v3.0': 64.5,
        'embed-multilingual-v3.0': 62.0,
        'voyage-2': 63.5,
        'voyage-large-2': 65.2,
        'all-MiniLM-L6-v2': 56.3,
        'all-mpnet-base-v2': 57.8,
        'bge-base-en-v1.5': 63.5,
        'bge-large-en-v1.5': 64.2,
    }

    # Pricing per 1M tokens
    PRICING = {
        'text-embedding-3-small': 0.02,
        'text-embedding-3-large': 0.13,
        'embed-english-v3.0': 0.10,
        'embed-multilingual-v3.0': 0.10,
        'voyage-2': 0.10,
        'voyage-large-2': 0.12,
        # Local models have zero API cost
        'all-MiniLM-L6-v2': 0.0,
        'all-mpnet-base-v2': 0.0,
        'bge-base-en-v1.5': 0.0,
        'bge-large-en-v1.5': 0.0,
    }

    def __init__(self, test_texts: Optional[List[str]] = None):
        """
        Initialize benchmark.

        Args:
            test_texts: Texts to use for benchmarking (default: sample texts)
        """
        if test_texts is None:
            self.test_texts = self._default_test_texts()
        else:
            self.test_texts = test_texts

        self.results: List[BenchmarkResult] = []

    def _default_test_texts(self) -> List[str]:
        """Generate default test texts for benchmarking."""
        return [
            "Machine learning is a subset of artificial intelligence.",
            "Natural language processing enables computers to understand human language.",
            "Deep learning uses neural networks with multiple layers.",
            "Computer vision allows machines to interpret visual information.",
            "Reinforcement learning trains agents through rewards and penalties.",
            "Supervised learning uses labeled data for training models.",
            "Unsupervised learning finds patterns in unlabeled data.",
            "Transfer learning adapts pre-trained models to new tasks.",
            "Generative AI creates new content like text and images.",
            "Large language models understand and generate human-like text.",
            "The quick brown fox jumps over the lazy dog.",
            "Paris is the capital and most populous city of France.",
            "Python is a high-level, interpreted programming language.",
            "Climate change refers to long-term shifts in temperatures.",
            "Quantum computing uses quantum phenomena for computation.",
        ] * 2  # 30 texts total for better statistics

    def benchmark_embedder(
        self,
        model_name: str,
        embed_func: Callable[[List[str]], List[List[float]]],
        provider: str,
        dimensions: int,
        batch_size: int = 10,
        warmup_runs: int = 2
    ) -> BenchmarkResult:
        """
        Benchmark a single embedder.

        Args:
            model_name: Name of the model
            embed_func: Function that takes texts and returns embeddings
            provider: Provider name ('openai', 'cohere', 'voyage', 'local')
            dimensions: Embedding dimensions
            batch_size: Batch size for embedding
            warmup_runs: Number of warmup runs before measurement

        Returns:
            BenchmarkResult with metrics
        """
        print(f"\nBenchmarking {model_name} ({provider})...")

        # Warmup
        for _ in range(warmup_runs):
            try:
                embed_func(self.test_texts[:batch_size])
            except Exception as e:
                print(f"  Warning: Warmup failed: {e}")

        # Measure latency
        latencies = []
        errors = 0

        for i in range(0, len(self.test_texts), batch_size):
            batch = self.test_texts[i:i + batch_size]

            start = time.time()
            try:
                embeddings = embed_func(batch)
                latency_ms = (time.time() - start) * 1000
                latencies.append(latency_ms)

                # Verify dimensions
                if len(embeddings) > 0 and len(embeddings[0]) != dimensions:
                    print(f"  Warning: Expected {dimensions} dims, got {len(embeddings[0])}")

            except Exception as e:
                print(f"  Error: {e}")
                errors += 1
                latencies.append(float('nan'))

        # Filter out errors
        valid_latencies = [l for l in latencies if not np.isnan(l)]

        if not valid_latencies:
            print(f"  Failed: No valid results")
            return BenchmarkResult(
                model_name=model_name,
                provider=provider,
                dimensions=dimensions,
                errors=errors,
                num_samples=0
            )

        # Calculate metrics
        total_time_sec = sum(valid_latencies) / 1000
        total_texts = len(self.test_texts) - (errors * batch_size)

        result = BenchmarkResult(
            model_name=model_name,
            provider=provider,
            dimensions=dimensions,
            latency_p50=np.percentile(valid_latencies, 50),
            latency_p95=np.percentile(valid_latencies, 95),
            latency_p99=np.percentile(valid_latencies, 99),
            latency_avg=np.mean(valid_latencies),
            throughput_texts_per_sec=total_texts / total_time_sec if total_time_sec > 0 else 0,
            cost_per_1m_tokens=self.PRICING.get(model_name, 0.0),
            mteb_score=self.MTEB_SCORES.get(model_name),
            batch_size=batch_size,
            num_samples=len(valid_latencies),
            errors=errors
        )

        print(f"  Latency p50: {result.latency_p50:.2f} ms")
        print(f"  Throughput: {result.throughput_texts_per_sec:.2f} texts/sec")
        print(f"  Cost per 1M tokens: ${result.cost_per_1m_tokens:.2f}")

        return result

    def benchmark_openai(self, model: str = 'text-embedding-3-small') -> Optional[BenchmarkResult]:
        """Benchmark OpenAI embedding model."""
        try:
            from openai import OpenAI

            if not os.getenv('OPENAI_API_KEY'):
                print(f"\nSkipping {model} (no OPENAI_API_KEY)")
                return None

            client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))

            dimensions = {
                'text-embedding-3-small': 1536,
                'text-embedding-3-large': 3072,
            }.get(model, 1536)

            def embed_func(texts: List[str]) -> List[List[float]]:
                response = client.embeddings.create(model=model, input=texts)
                return [d.embedding for d in response.data]

            return self.benchmark_embedder(
                model_name=model,
                embed_func=embed_func,
                provider='openai',
                dimensions=dimensions,
                batch_size=20
            )

        except ImportError:
            print(f"\nSkipping {model} (openai not installed)")
            return None
        except Exception as e:
            print(f"\nError benchmarking {model}: {e}")
            return None

    def benchmark_cohere(self, model: str = 'embed-english-v3.0') -> Optional[BenchmarkResult]:
        """Benchmark Cohere embedding model."""
        try:
            import cohere

            if not os.getenv('COHERE_API_KEY'):
                print(f"\nSkipping {model} (no COHERE_API_KEY)")
                return None

            client = cohere.Client(api_key=os.getenv('COHERE_API_KEY'))

            dimensions = 1024  # Cohere v3 models

            def embed_func(texts: List[str]) -> List[List[float]]:
                response = client.embed(texts=texts, model=model, input_type='search_document')
                return response.embeddings

            return self.benchmark_embedder(
                model_name=model,
                embed_func=embed_func,
                provider='cohere',
                dimensions=dimensions,
                batch_size=20
            )

        except ImportError:
            print(f"\nSkipping {model} (cohere not installed)")
            return None
        except Exception as e:
            print(f"\nError benchmarking {model}: {e}")
            return None

    def benchmark_voyage(self, model: str = 'voyage-2') -> Optional[BenchmarkResult]:
        """Benchmark Voyage AI embedding model."""
        try:
            import voyageai

            if not os.getenv('VOYAGE_API_KEY'):
                print(f"\nSkipping {model} (no VOYAGE_API_KEY)")
                return None

            client = voyageai.Client(api_key=os.getenv('VOYAGE_API_KEY'))

            dimensions = {
                'voyage-2': 1024,
                'voyage-large-2': 1536,
            }.get(model, 1024)

            def embed_func(texts: List[str]) -> List[List[float]]:
                response = client.embed(texts=texts, model=model)
                return response.embeddings

            return self.benchmark_embedder(
                model_name=model,
                embed_func=embed_func,
                provider='voyage',
                dimensions=dimensions,
                batch_size=20
            )

        except ImportError:
            print(f"\nSkipping {model} (voyageai not installed)")
            return None
        except Exception as e:
            print(f"\nError benchmarking {model}: {e}")
            return None

    def benchmark_local(self, model: str = 'all-MiniLM-L6-v2') -> Optional[BenchmarkResult]:
        """Benchmark local sentence-transformers model."""
        try:
            from sentence_transformers import SentenceTransformer

            print(f"\nLoading local model {model}...")
            encoder = SentenceTransformer(model)

            dimensions = {
                'all-MiniLM-L6-v2': 384,
                'all-mpnet-base-v2': 768,
                'bge-base-en-v1.5': 768,
                'bge-large-en-v1.5': 1024,
            }.get(model, 384)

            def embed_func(texts: List[str]) -> List[List[float]]:
                embeddings = encoder.encode(texts, convert_to_numpy=True)
                return embeddings.tolist()

            return self.benchmark_embedder(
                model_name=model,
                embed_func=embed_func,
                provider='local',
                dimensions=dimensions,
                batch_size=16
            )

        except ImportError:
            print(f"\nSkipping {model} (sentence-transformers not installed)")
            return None
        except Exception as e:
            print(f"\nError benchmarking {model}: {e}")
            return None

    def run_benchmarks(self, models: Optional[Dict[str, List[str]]] = None) -> List[BenchmarkResult]:
        """
        Run benchmarks for multiple models.

        Args:
            models: Dictionary mapping provider to list of model names
                   If None, uses default model set

        Returns:
            List of BenchmarkResults
        """
        if models is None:
            models = {
                'openai': ['text-embedding-3-small', 'text-embedding-3-large'],
                'cohere': ['embed-english-v3.0'],
                'voyage': ['voyage-2'],
                'local': ['all-MiniLM-L6-v2', 'bge-base-en-v1.5']
            }

        self.results = []

        # Benchmark OpenAI models
        for model in models.get('openai', []):
            result = self.benchmark_openai(model)
            if result:
                self.results.append(result)

        # Benchmark Cohere models
        for model in models.get('cohere', []):
            result = self.benchmark_cohere(model)
            if result:
                self.results.append(result)

        # Benchmark Voyage models
        for model in models.get('voyage', []):
            result = self.benchmark_voyage(model)
            if result:
                self.results.append(result)

        # Benchmark local models
        for model in models.get('local', []):
            result = self.benchmark_local(model)
            if result:
                self.results.append(result)

        return self.results

    def print_summary(self):
        """Print benchmark summary table."""
        if not self.results:
            print("\nNo benchmark results available.")
            return

        print("\n" + "=" * 100)
        print("EMBEDDING MODEL BENCHMARK SUMMARY")
        print("=" * 100)

        # Header
        print(f"\n{'Model':<30} {'Provider':<10} {'Dims':<6} {'p50 (ms)':<10} {'p95 (ms)':<10} "
              f"{'Throughput':<15} {'Cost/1M':<10} {'MTEB':<8}")
        print("-" * 100)

        # Results
        for r in sorted(self.results, key=lambda x: x.latency_p50):
            throughput_str = f"{r.throughput_texts_per_sec:.1f} tx/s"
            cost_str = f"${r.cost_per_1m_tokens:.3f}" if r.cost_per_1m_tokens > 0 else "Free"
            mteb_str = f"{r.mteb_score:.1f}" if r.mteb_score else "N/A"

            print(f"{r.model_name:<30} {r.provider:<10} {r.dimensions:<6} "
                  f"{r.latency_p50:<10.2f} {r.latency_p95:<10.2f} "
                  f"{throughput_str:<15} {cost_str:<10} {mteb_str:<8}")

        print("\n" + "=" * 100)

    def visualize_results(self, save_path: Optional[str] = None):
        """
        Create visualization of benchmark results.

        Args:
            save_path: Optional path to save figure
        """
        if not self.results:
            print("\nNo results to visualize.")
            return

        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
        fig.suptitle('Embedding Model Benchmark Results', fontsize=16, fontweight='bold')

        # Extract data
        models = [r.model_name for r in self.results]
        providers = [r.provider for r in self.results]
        colors = {'openai': '#10a37f', 'cohere': '#39594D', 'voyage': '#4B8BBE', 'local': '#646464'}
        model_colors = [colors.get(p, '#000000') for p in providers]

        # 1. Latency comparison
        ax = axes[0, 0]
        latencies_p50 = [r.latency_p50 for r in self.results]
        latencies_p95 = [r.latency_p95 for r in self.results]

        x = np.arange(len(models))
        width = 0.35
        ax.bar(x - width/2, latencies_p50, width, label='p50', color=model_colors, alpha=0.8)
        ax.bar(x + width/2, latencies_p95, width, label='p95', color=model_colors, alpha=0.5)

        ax.set_ylabel('Latency (ms)')
        ax.set_title('Latency by Model')
        ax.set_xticks(x)
        ax.set_xticklabels([m.replace('text-embedding-', '').replace('embed-', '')[:15] for m in models],
                          rotation=45, ha='right')
        ax.legend()
        ax.grid(axis='y', alpha=0.3)

        # 2. Throughput comparison
        ax = axes[0, 1]
        throughputs = [r.throughput_texts_per_sec for r in self.results]
        ax.bar(models, throughputs, color=model_colors, alpha=0.8)
        ax.set_ylabel('Throughput (texts/sec)')
        ax.set_title('Throughput by Model')
        ax.set_xticklabels([m.replace('text-embedding-', '').replace('embed-', '')[:15] for m in models],
                          rotation=45, ha='right')
        ax.grid(axis='y', alpha=0.3)

        # 3. Cost vs Quality (scatter)
        ax = axes[1, 0]
        costs = [r.cost_per_1k_texts() for r in self.results]
        qualities = [r.mteb_score if r.mteb_score else 0 for r in self.results]

        for i, r in enumerate(self.results):
            if r.mteb_score:
                ax.scatter(costs[i], qualities[i], s=200, c=model_colors[i], alpha=0.7)
                # Label with shortened model name
                label = r.model_name.replace('text-embedding-', '').replace('embed-', '')[:10]
                ax.annotate(label, (costs[i], qualities[i]),
                          xytext=(5, 5), textcoords='offset points', fontsize=8)

        ax.set_xlabel('Cost per 1K texts ($)')
        ax.set_ylabel('MTEB Score')
        ax.set_title('Quality vs Cost Trade-off')
        ax.grid(alpha=0.3)

        # 4. Dimensions vs Latency (scatter)
        ax = axes[1, 1]
        dimensions = [r.dimensions for r in self.results]

        ax.scatter(dimensions, latencies_p50, s=200, c=model_colors, alpha=0.7)
        for i, r in enumerate(self.results):
            label = r.model_name.replace('text-embedding-', '').replace('embed-', '')[:10]
            ax.annotate(label, (dimensions[i], latencies_p50[i]),
                       xytext=(5, 5), textcoords='offset points', fontsize=8)

        ax.set_xlabel('Embedding Dimensions')
        ax.set_ylabel('Latency p50 (ms)')
        ax.set_title('Dimensions vs Latency')
        ax.grid(alpha=0.3)

        plt.tight_layout()

        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches='tight')
            print(f"\nVisualization saved to {save_path}")
        else:
            plt.show()


def main():
    """Run embedding model benchmarks."""

    print("=" * 100)
    print("EMBEDDING MODEL BENCHMARK")
    print("=" * 100)
    print("\nThis benchmark will compare:")
    print("  - OpenAI: text-embedding-3-small, text-embedding-3-large")
    print("  - Cohere: embed-english-v3.0")
    print("  - Voyage AI: voyage-2")
    print("  - Local: all-MiniLM-L6-v2, bge-base-en-v1.5")
    print("\nMetrics: latency, throughput, cost, quality (MTEB score)")
    print("\nNote: API keys required for cloud providers (set as env vars)")
    print("=" * 100)

    # Initialize benchmark
    benchmark = EmbedderBenchmark()

    # Run benchmarks
    print("\nStarting benchmarks...")
    results = benchmark.run_benchmarks()

    # Print summary
    benchmark.print_summary()

    # Print recommendations
    print("\n" + "=" * 100)
    print("RECOMMENDATIONS")
    print("=" * 100)

    if results:
        # Best by latency (local)
        local_results = [r for r in results if r.provider == 'local']
        if local_results:
            fastest_local = min(local_results, key=lambda x: x.latency_p50)
            print(f"\nFastest Local Model: {fastest_local.model_name}")
            print(f"  - Latency p50: {fastest_local.latency_p50:.2f} ms")
            print(f"  - Dimensions: {fastest_local.dimensions}")
            print(f"  - MTEB Score: {fastest_local.mteb_score or 'N/A'}")
            print(f"  - Use case: High volume, tight budgets, data privacy")

        # Best API model (quality/cost balance)
        api_results = [r for r in results if r.provider != 'local' and r.mteb_score]
        if api_results:
            best_value = max(api_results, key=lambda x: x.quality_per_dollar())
            print(f"\nBest Value API Model: {best_value.model_name}")
            print(f"  - MTEB Score: {best_value.mteb_score:.1f}")
            print(f"  - Cost per 1K texts: ${best_value.cost_per_1k_texts():.4f}")
            print(f"  - Quality per dollar: {best_value.quality_per_dollar():.2f}")
            print(f"  - Use case: Production RAG systems")

        # Highest quality
        if api_results:
            best_quality = max(api_results, key=lambda x: x.mteb_score or 0)
            print(f"\nHighest Quality: {best_quality.model_name}")
            print(f"  - MTEB Score: {best_quality.mteb_score:.1f}")
            print(f"  - Dimensions: {best_quality.dimensions}")
            print(f"  - Cost per 1K texts: ${best_quality.cost_per_1k_texts():.4f}")
            print(f"  - Use case: Premium applications, accuracy-critical")

    print("\n" + "=" * 100)

    # Visualize
    print("\nGenerating visualization...")
    benchmark.visualize_results(save_path='benchmark_results.png')

    print("\nBenchmark complete!")


if __name__ == "__main__":
    main()

"""
Local Embedding with sentence-transformers

Demonstrates zero-cost local embedding with GPU acceleration.

Dependencies:
    pip install sentence-transformers torch

Usage:
    python local_embedder.py

Notes:
    - Automatically detects GPU (CUDA, MPS for Apple Silicon)
    - Falls back to CPU if no GPU available
    - First run downloads model (~90MB for all-MiniLM-L6-v2)
"""

from sentence_transformers import SentenceTransformer
from typing import List
import torch
import time


class LocalEmbedder:
    """Local embedding model with GPU acceleration."""

    def __init__(
        self,
        model_name: str = "all-MiniLM-L6-v2",
        device: str = None,
        batch_size: int = 32
    ):
        """
        Initialize local embedder.

        Args:
            model_name: Model from sentence-transformers Hub
                Available models:
                - 'all-MiniLM-L6-v2': 384 dims, fast (22M params)
                - 'BAAI/bge-base-en-v1.5': 768 dims, SOTA quality (109M params)
                - 'BAAI/bge-large-en-v1.5': 1024 dims, max quality (335M params)
                - 'paraphrase-multilingual-MiniLM-L12-v2': 384 dims, 50+ languages
            device: 'cuda', 'mps' (Apple Silicon), 'cpu', or None (auto-detect)
            batch_size: Batch size for encoding (tune based on GPU memory)
        """
        if device is None:
            device = self._auto_detect_device()

        self.model = SentenceTransformer(model_name, device=device)
        self.model_name = model_name
        self.device = device
        self.batch_size = batch_size

        print(f"Loaded {model_name} on {device}")
        print(f"Embedding dimension: {self.get_embedding_dim()}")

    def _auto_detect_device(self) -> str:
        """Auto-detect best available device."""
        if torch.cuda.is_available():
            return 'cuda'
        elif torch.backends.mps.is_available():
            return 'mps'
        else:
            return 'cpu'

    def embed_single(self, text: str) -> List[float]:
        """
        Embed single text.

        Args:
            text: Text to embed

        Returns:
            Embedding vector (list of floats)
        """
        embedding = self.model.encode(text, convert_to_numpy=True)
        return embedding.tolist()

    def embed_batch(
        self,
        texts: List[str],
        show_progress: bool = False,
        normalize_embeddings: bool = False
    ) -> List[List[float]]:
        """
        Embed batch of texts with optimized batching.

        Args:
            texts: List of texts to embed
            show_progress: Show progress bar (useful for large batches)
            normalize_embeddings: Normalize embeddings to unit length (for cosine similarity)

        Returns:
            List of embedding vectors (each is a list of floats)
        """
        embeddings = self.model.encode(
            texts,
            batch_size=self.batch_size,
            show_progress_bar=show_progress,
            convert_to_numpy=True,
            normalize_embeddings=normalize_embeddings
        )
        return embeddings.tolist()

    def get_embedding_dim(self) -> int:
        """Return embedding dimensionality."""
        return self.model.get_sentence_embedding_dimension()

    def get_max_seq_length(self) -> int:
        """Return maximum sequence length (in tokens)."""
        return self.model.get_max_seq_length()


def benchmark_throughput(embedder: LocalEmbedder, num_texts: int = 1000):
    """
    Benchmark embedding throughput.

    Args:
        embedder: LocalEmbedder instance
        num_texts: Number of texts to embed for benchmark
    """
    print(f"\nBenchmarking throughput ({num_texts} texts)...")

    # Generate sample texts
    texts = [
        f"This is sample text number {i} for benchmarking embedding throughput."
        for i in range(num_texts)
    ]

    # Benchmark
    start = time.time()
    embeddings = embedder.embed_batch(texts, show_progress=True)
    elapsed = time.time() - start

    throughput = num_texts / elapsed
    avg_latency = (elapsed / num_texts) * 1000  # ms per text

    print(f"\nBenchmark Results:")
    print(f"  Total texts: {num_texts}")
    print(f"  Total time: {elapsed:.2f}s")
    print(f"  Throughput: {throughput:.2f} texts/sec")
    print(f"  Avg latency: {avg_latency:.2f}ms per text")
    print(f"  Device: {embedder.device}")


def compare_models():
    """Compare different embedding models."""
    print("\nComparing different models...")
    print("=" * 60)

    models = [
        ("all-MiniLM-L6-v2", "Fast, small (384 dims)"),
        # Uncomment to test (requires download):
        # ("BAAI/bge-base-en-v1.5", "High quality (768 dims)"),
    ]

    sample_text = "What is machine learning?"

    for model_name, description in models:
        print(f"\nModel: {model_name}")
        print(f"Description: {description}")

        embedder = LocalEmbedder(model_name=model_name)

        # Embed sample
        start = time.time()
        embedding = embedder.embed_single(sample_text)
        latency = (time.time() - start) * 1000

        print(f"  Dimension: {len(embedding)}")
        print(f"  Latency: {latency:.2f}ms")
        print(f"  First 5 values: {[round(v, 4) for v in embedding[:5]]}")


def main():
    """Example usage of LocalEmbedder."""

    print("Local Embedding with sentence-transformers Demo")
    print("=" * 60)

    # Initialize embedder (auto-detects GPU)
    embedder = LocalEmbedder(model_name="all-MiniLM-L6-v2")

    print(f"\nDevice: {embedder.device}")
    print(f"Max sequence length: {embedder.get_max_seq_length()} tokens")

    # Example 1: Single text embedding
    print("\n1. Single text embedding:")
    text = "What is the capital of France?"
    embedding = embedder.embed_single(text)
    print(f"   Text: {text}")
    print(f"   Embedding dimension: {len(embedding)}")
    print(f"   First 5 values: {[round(v, 4) for v in embedding[:5]]}")

    # Example 2: Batch embedding
    print("\n2. Batch embedding:")
    texts = [
        "Paris is the capital of France.",
        "London is the capital of the UK.",
        "Berlin is the capital of Germany.",
        "Madrid is the capital of Spain.",
        "Rome is the capital of Italy."
    ]

    embeddings = embedder.embed_batch(texts)
    print(f"   Embedded {len(embeddings)} texts")
    print(f"   All embeddings have dimension: {len(embeddings[0])}")

    # Example 3: Semantic similarity
    print("\n3. Semantic similarity:")
    query = "French capital city"
    query_embedding = embedder.embed_single(query)

    # Calculate cosine similarity with each text
    def cosine_similarity(a, b):
        import numpy as np
        return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

    print(f"   Query: {query}")
    for i, text in enumerate(texts):
        similarity = cosine_similarity(query_embedding, embeddings[i])
        print(f"   {text}: {similarity:.4f}")

    # Example 4: Normalized embeddings (for faster cosine similarity)
    print("\n4. Normalized embeddings (unit length):")
    normalized_embeddings = embedder.embed_batch(texts, normalize_embeddings=True)

    # With normalized embeddings, cosine similarity = dot product
    import numpy as np
    query_norm = embedder.embed_batch([query], normalize_embeddings=True)[0]

    print(f"   Query: {query}")
    for i, text in enumerate(texts):
        similarity = np.dot(query_norm, normalized_embeddings[i])
        print(f"   {text}: {similarity:.4f}")

    # Example 5: Throughput benchmark
    if embedder.device in ['cuda', 'mps']:
        benchmark_throughput(embedder, num_texts=1000)
    else:
        print("\n5. Throughput benchmark:")
        print("   Skipping (CPU inference is slow, use GPU for benchmarking)")

    # Example 6: Cost comparison
    print("\n6. Cost comparison (local vs. API):")
    print("   Local model (all-MiniLM-L6-v2):")
    print("   - API cost: $0 (free, runs locally)")
    print("   - Infrastructure: GPU recommended (~$150-400/month)")
    print("   - Throughput: 5,000-10,000 texts/sec (GPU)")
    print("\n   OpenAI API (text-embedding-3-small):")
    print("   - API cost: $0.02 per 1M tokens (~$0.50/month for 10M chars)")
    print("   - Infrastructure: None (managed)")
    print("   - Throughput: 1,000-5,000 texts/min (rate-limited)")
    print("\n   Break-even: Local is cost-effective above ~1M embeddings/month")

    print("\n" + "=" * 60)
    print("Demo complete!")


if __name__ == "__main__":
    main()

"""
OpenAI Embedding with Redis Caching

Demonstrates content-addressable caching to reduce API costs by 80-90%.

Dependencies:
    pip install openai redis

Usage:
    python openai_cached.py

Environment Variables:
    OPENAI_API_KEY: Your OpenAI API key
    REDIS_HOST: Redis server host (default: localhost)
    REDIS_PORT: Redis server port (default: 6379)
"""

import hashlib
import json
import os
from typing import List, Dict, Optional
from openai import OpenAI
import redis


class CachedEmbedder:
    """OpenAI embedder with Redis content-addressable caching."""

    def __init__(
        self,
        model: str = "text-embedding-3-small",
        redis_host: str = "localhost",
        redis_port: int = 6379,
        cache_ttl: int = 86400 * 30  # 30 days
    ):
        """
        Initialize cached embedder.

        Args:
            model: OpenAI embedding model name
            redis_host: Redis server hostname
            redis_port: Redis server port
            cache_ttl: Cache time-to-live in seconds
        """
        self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
        self.model = model
        self.redis_client = redis.Redis(
            host=redis_host,
            port=redis_port,
            decode_responses=False
        )
        self.cache_ttl = cache_ttl

        # Statistics
        self.cache_hits = 0
        self.cache_misses = 0
        self.api_calls = 0

    def _cache_key(self, text: str) -> str:
        """
        Generate content-addressable cache key.

        Args:
            text: Input text

        Returns:
            Cache key (SHA-256 hash)
        """
        content_hash = hashlib.sha256(text.encode('utf-8')).hexdigest()
        return f"embed:{self.model}:{content_hash}"

    def embed_single(self, text: str) -> List[float]:
        """
        Embed single text with caching.

        Args:
            text: Text to embed

        Returns:
            Embedding vector (list of floats)
        """
        cache_key = self._cache_key(text)

        # Check cache
        cached = self.redis_client.get(cache_key)
        if cached:
            self.cache_hits += 1
            return json.loads(cached)

        # Generate embedding via API
        self.cache_misses += 1
        self.api_calls += 1

        try:
            response = self.client.embeddings.create(
                model=self.model,
                input=text
            )
            embedding = response.data[0].embedding

            # Cache result
            self.redis_client.setex(
                cache_key,
                self.cache_ttl,
                json.dumps(embedding)
            )

            return embedding

        except Exception as e:
            print(f"Error embedding text: {e}")
            raise

    def embed_batch(self, texts: List[str], batch_size: int = 2048) -> List[List[float]]:
        """
        Embed batch of texts with caching (up to 2,048 texts per API call).

        Args:
            texts: List of texts to embed
            batch_size: Maximum batch size for OpenAI API

        Returns:
            List of embedding vectors
        """
        if len(texts) > batch_size:
            # Split into multiple batches
            embeddings = []
            for i in range(0, len(texts), batch_size):
                batch = texts[i:i + batch_size]
                embeddings.extend(self.embed_batch(batch, batch_size))
            return embeddings

        # Prepare result array and track uncached texts
        embeddings = []
        uncached_texts = []
        uncached_indices = []

        # Check cache for all texts
        for i, text in enumerate(texts):
            cache_key = self._cache_key(text)
            cached = self.redis_client.get(cache_key)

            if cached:
                self.cache_hits += 1
                embeddings.append(json.loads(cached))
            else:
                embeddings.append(None)
                uncached_texts.append(text)
                uncached_indices.append(i)

        # Generate embeddings for uncached texts
        if uncached_texts:
            self.cache_misses += len(uncached_texts)
            self.api_calls += 1

            try:
                response = self.client.embeddings.create(
                    model=self.model,
                    input=uncached_texts
                )

                # Insert into results and cache
                for idx, data in zip(uncached_indices, response.data):
                    embedding = data.embedding
                    embeddings[idx] = embedding

                    # Cache result
                    cache_key = self._cache_key(texts[idx])
                    self.redis_client.setex(
                        cache_key,
                        self.cache_ttl,
                        json.dumps(embedding)
                    )

            except Exception as e:
                print(f"Error embedding batch: {e}")
                raise

        return embeddings

    def get_cache_stats(self) -> Dict[str, float]:
        """
        Get cache hit/miss statistics.

        Returns:
            Dictionary with cache statistics
        """
        total = self.cache_hits + self.cache_misses
        hit_rate = (self.cache_hits / total * 100) if total > 0 else 0

        return {
            'cache_hits': self.cache_hits,
            'cache_misses': self.cache_misses,
            'total_requests': total,
            'hit_rate_pct': round(hit_rate, 2),
            'api_calls': self.api_calls,
            'api_savings_pct': round((1 - self.api_calls / max(1, total)) * 100, 2)
        }

    def clear_cache(self) -> int:
        """
        Clear all cached embeddings for this model.

        Returns:
            Number of keys deleted
        """
        pattern = f"embed:{self.model}:*"
        keys = self.redis_client.keys(pattern)
        if keys:
            return self.redis_client.delete(*keys)
        return 0


def main():
    """Example usage of CachedEmbedder."""

    print("OpenAI Embedding with Redis Caching Demo")
    print("=" * 60)

    # Initialize embedder
    embedder = CachedEmbedder(model="text-embedding-3-small")

    # Example 1: Single text embedding
    print("\n1. Single text embedding:")
    text = "What is the capital of France?"
    embedding = embedder.embed_single(text)
    print(f"   Text: {text}")
    print(f"   Embedding dimension: {len(embedding)}")
    print(f"   First 5 values: {embedding[:5]}")

    # Example 2: Embed same text again (cache hit)
    print("\n2. Embed same text again (should be cached):")
    embedding = embedder.embed_single(text)
    stats = embedder.get_cache_stats()
    print(f"   Cache hit rate: {stats['hit_rate_pct']}%")
    print(f"   API calls: {stats['api_calls']}")

    # Example 3: Batch embedding
    print("\n3. Batch embedding:")
    texts = [
        "Paris is the capital of France.",
        "London is the capital of the UK.",
        "Berlin is the capital of Germany.",
        "What is the capital of France?",  # Duplicate (cached)
        "Madrid is the capital of Spain."
    ]

    embeddings = embedder.embed_batch(texts)
    print(f"   Embedded {len(embeddings)} texts")
    print(f"   All embeddings have dimension: {len(embeddings[0])}")

    # Example 4: Cache statistics
    print("\n4. Cache statistics:")
    stats = embedder.get_cache_stats()
    print(f"   Total requests: {stats['total_requests']}")
    print(f"   Cache hits: {stats['cache_hits']}")
    print(f"   Cache misses: {stats['cache_misses']}")
    print(f"   Hit rate: {stats['hit_rate_pct']}%")
    print(f"   API calls: {stats['api_calls']}")
    print(f"   API call savings: {stats['api_savings_pct']}%")

    # Example 5: Cost estimation
    print("\n5. Cost estimation:")
    avg_tokens = 25  # Estimated average tokens per text
    total_tokens_without_cache = stats['total_requests'] * avg_tokens
    total_tokens_with_cache = stats['api_calls'] * avg_tokens
    cost_per_1m_tokens = 0.02  # text-embedding-3-small pricing

    cost_without_cache = (total_tokens_without_cache / 1_000_000) * cost_per_1m_tokens
    cost_with_cache = (total_tokens_with_cache / 1_000_000) * cost_per_1m_tokens
    savings = cost_without_cache - cost_with_cache

    print(f"   Without caching: ${cost_without_cache:.6f}")
    print(f"   With caching: ${cost_with_cache:.6f}")
    print(f"   Savings: ${savings:.6f} ({stats['api_savings_pct']}%)")

    print("\n" + "=" * 60)
    print("Demo complete!")


if __name__ == "__main__":
    main()

"""
Embedding Pipeline Performance Monitoring

Demonstrates tracking latency, throughput, cost, and cache efficiency.

Dependencies:
    pip install numpy

Usage:
    python performance_monitor.py
"""

import time
import numpy as np
from typing import List, Dict, Callable
from dataclasses import dataclass, field
from datetime import datetime


@dataclass
class EmbeddingMetrics:
    """Metrics for embedding operations."""

    total_texts: int = 0
    total_embeddings: int = 0
    total_api_calls: int = 0
    total_cache_hits: int = 0
    total_latency_ms: float = 0.0
    total_cost_usd: float = 0.0
    errors: int = 0

    start_time: datetime = field(default_factory=datetime.now)

    def add_operation(
        self,
        num_texts: int,
        latency_ms: float,
        cache_hits: int = 0,
        api_calls: int = 1,
        cost_usd: float = 0.0,
        error: bool = False
    ):
        """Record an embedding operation."""
        self.total_texts += num_texts
        self.total_embeddings += num_texts
        self.total_api_calls += api_calls
        self.total_cache_hits += cache_hits
        self.total_latency_ms += latency_ms
        self.total_cost_usd += cost_usd
        if error:
            self.errors += 1

    def get_summary(self) -> Dict:
        """Get performance summary."""
        elapsed_seconds = (datetime.now() - self.start_time).total_seconds()

        cache_hit_rate = (
            (self.total_cache_hits / self.total_texts * 100)
            if self.total_texts > 0 else 0
        )

        throughput = (
            self.total_texts / elapsed_seconds
            if elapsed_seconds > 0 else 0
        )

        avg_latency = (
            self.total_latency_ms / self.total_api_calls
            if self.total_api_calls > 0 else 0
        )

        error_rate = (
            (self.errors / self.total_api_calls * 100)
            if self.total_api_calls > 0 else 0
        )

        return {
            'total_texts': self.total_texts,
            'total_api_calls': self.total_api_calls,
            'cache_hit_rate_pct': round(cache_hit_rate, 2),
            'throughput_texts_per_sec': round(throughput, 2),
            'avg_latency_ms': round(avg_latency, 2),
            'total_cost_usd': round(self.total_cost_usd, 4),
            'error_rate_pct': round(error_rate, 2),
            'elapsed_seconds': round(elapsed_seconds, 2)
        }


class LatencyTracker:
    """Track latency percentiles."""

    def __init__(self):
        self.latencies = []

    def record(self, latency_ms: float):
        """Record a latency measurement."""
        self.latencies.append(latency_ms)

    def percentiles(self) -> Dict:
        """Get latency percentiles."""
        if not self.latencies:
            return {}

        return {
            'p50': round(np.percentile(self.latencies, 50), 2),
            'p95': round(np.percentile(self.latencies, 95), 2),
            'p99': round(np.percentile(self.latencies, 99), 2),
            'avg': round(np.mean(self.latencies), 2),
            'min': round(np.min(self.latencies), 2),
            'max': round(np.max(self.latencies), 2)
        }

    def reset(self):
        """Reset latency measurements."""
        self.latencies = []


class CostTracker:
    """Track embedding costs."""

    PRICING = {
        'text-embedding-3-small': 0.00002,  # per 1K tokens
        'text-embedding-3-large': 0.00013,
        'embed-english-v3.0': 0.0001,
        'local': 0.0,  # Local models have zero API cost
    }

    def __init__(self, model: str = 'text-embedding-3-small'):
        self.model = model
        self.total_tokens = 0
        self.total_cost = 0.0
        self.calls = 0

    def _estimate_tokens(self, text: str) -> int:
        """Rough token estimation (4 chars per token)."""
        return len(text) // 4

    def record_embedding(self, text: str, from_cache: bool = False) -> Dict:
        """
        Record embedding cost.

        Args:
            text: Input text
            from_cache: Whether this was served from cache (no cost)

        Returns:
            Dictionary with tokens and cost
        """
        tokens = self._estimate_tokens(text)
        cost = 0.0

        if not from_cache:
            cost = (tokens / 1000) * self.PRICING.get(self.model, 0)
            self.total_tokens += tokens
            self.total_cost += cost
            self.calls += 1

        return {
            'tokens': tokens,
            'cost_usd': cost
        }

    def summary(self) -> Dict:
        """Get cost summary."""
        avg_tokens = self.total_tokens / max(1, self.calls)
        avg_cost = self.total_cost / max(1, self.calls)

        return {
            'total_tokens': self.total_tokens,
            'total_cost_usd': round(self.total_cost, 4),
            'avg_tokens_per_text': round(avg_tokens, 2),
            'avg_cost_per_text': round(avg_cost, 6),
            'total_calls': self.calls
        }


class MonitoredEmbedder:
    """Wrapper for any embedder with performance monitoring."""

    def __init__(
        self,
        embedder: any,
        model: str = 'text-embedding-3-small'
    ):
        """
        Initialize monitored embedder.

        Args:
            embedder: Underlying embedder instance
            model: Model name for cost tracking
        """
        self.embedder = embedder
        self.metrics = EmbeddingMetrics()
        self.latency_tracker = LatencyTracker()
        self.cost_tracker = CostTracker(model)

    def embed_single(self, text: str) -> List[float]:
        """Embed with monitoring."""
        start = time.time()

        # Check if embedder has cache
        from_cache = False
        if hasattr(self.embedder, '_cache_key') and hasattr(self.embedder, 'redis_client'):
            cache_key = self.embedder._cache_key(text)
            from_cache = self.embedder.redis_client.exists(cache_key)

        # Embed
        try:
            embedding = self.embedder.embed_single(text)

            # Record metrics
            latency_ms = (time.time() - start) * 1000
            self.latency_tracker.record(latency_ms)

            cost_info = self.cost_tracker.record_embedding(text, from_cache)

            self.metrics.add_operation(
                num_texts=1,
                latency_ms=latency_ms,
                cache_hits=1 if from_cache else 0,
                api_calls=0 if from_cache else 1,
                cost_usd=cost_info['cost_usd']
            )

            return embedding

        except Exception as e:
            self.metrics.add_operation(
                num_texts=1,
                latency_ms=(time.time() - start) * 1000,
                error=True
            )
            raise

    def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """Embed batch with monitoring."""
        start = time.time()

        # Check cache hits
        cache_hits = 0
        if hasattr(self.embedder, '_cache_key') and hasattr(self.embedder, 'redis_client'):
            for text in texts:
                cache_key = self.embedder._cache_key(text)
                if self.embedder.redis_client.exists(cache_key):
                    cache_hits += 1

        # Embed
        try:
            embeddings = self.embedder.embed_batch(texts)

            # Record metrics
            latency_ms = (time.time() - start) * 1000
            self.latency_tracker.record(latency_ms)

            # Calculate costs
            total_cost = 0.0
            for text in texts:
                from_cache = False
                if hasattr(self.embedder, '_cache_key') and hasattr(self.embedder, 'redis_client'):
                    cache_key = self.embedder._cache_key(text)
                    from_cache = self.embedder.redis_client.exists(cache_key)

                cost_info = self.cost_tracker.record_embedding(text, from_cache)
                total_cost += cost_info['cost_usd']

            api_calls = 1 if cache_hits < len(texts) else 0

            self.metrics.add_operation(
                num_texts=len(texts),
                latency_ms=latency_ms,
                cache_hits=cache_hits,
                api_calls=api_calls,
                cost_usd=total_cost
            )

            return embeddings

        except Exception as e:
            self.metrics.add_operation(
                num_texts=len(texts),
                latency_ms=(time.time() - start) * 1000,
                error=True
            )
            raise

    def get_metrics(self) -> Dict:
        """Get comprehensive performance metrics."""
        summary = self.metrics.get_summary()
        latency = self.latency_tracker.percentiles()
        cost = self.cost_tracker.summary()

        return {
            **summary,
            'latency': latency,
            'cost': cost
        }

    def print_dashboard(self):
        """Print a comprehensive dashboard."""
        metrics = self.get_metrics()

        print("=" * 60)
        print("EMBEDDING PIPELINE DASHBOARD")
        print("=" * 60)

        print("\nPERFORMANCE:")
        print(f"  Throughput:     {metrics['throughput_texts_per_sec']:.2f} texts/sec")
        print(f"  Avg Latency:    {metrics['avg_latency_ms']:.2f} ms")
        if metrics.get('latency'):
            print(f"  Latency p50:    {metrics['latency']['p50']:.2f} ms")
            print(f"  Latency p95:    {metrics['latency']['p95']:.2f} ms")
            print(f"  Latency p99:    {metrics['latency']['p99']:.2f} ms")
        print(f"  Error Rate:     {metrics['error_rate_pct']:.2f}%")

        print("\nCACHE:")
        print(f"  Hit Rate:       {metrics['cache_hit_rate_pct']:.2f}%")
        print(f"  Cache Hits:     {metrics['total_cache_hits']:,}")
        print(f"  Cache Misses:   {metrics['total_texts'] - metrics['total_cache_hits']:,}")

        print("\nCOST:")
        if metrics.get('cost'):
            print(f"  Total Cost:     ${metrics['cost']['total_cost_usd']:.4f}")
            print(f"  Total Tokens:   {metrics['cost']['total_tokens']:,}")
            print(f"  Avg Cost/Text:  ${metrics['cost']['avg_cost_per_text']:.6f}")

        print("\nOVERALL:")
        print(f"  Total Texts:    {metrics['total_texts']:,}")
        print(f"  API Calls:      {metrics['total_api_calls']:,}")
        print(f"  Elapsed Time:   {metrics['elapsed_seconds']:.2f}s")

        print("=" * 60)


# Mock embedder for demonstration
class MockEmbedder:
    """Mock embedder for testing monitoring."""

    def embed_single(self, text: str) -> List[float]:
        time.sleep(0.01)  # Simulate API latency
        return [0.1] * 384

    def embed_batch(self, texts: List[str]) -> List[List[float]]:
        time.sleep(0.05)  # Simulate batch API latency
        return [[0.1] * 384 for _ in texts]


def main():
    """Example usage of performance monitoring."""

    print("Embedding Pipeline Performance Monitoring Demo")
    print("=" * 60)

    # Create mock embedder
    embedder = MockEmbedder()
    monitored = MonitoredEmbedder(embedder, model='text-embedding-3-small')

    # Example 1: Single text embeddings
    print("\n1. Embedding single texts...")
    for i in range(10):
        text = f"Sample text number {i}"
        monitored.embed_single(text)

    print("   Embedded 10 texts")

    # Example 2: Batch embedding
    print("\n2. Embedding batch...")
    texts = [f"Batch text {i}" for i in range(50)]
    monitored.embed_batch(texts)

    print("   Embedded 50 texts in batch")

    # Example 3: Show metrics
    print("\n3. Performance Metrics:")
    metrics = monitored.get_metrics()

    print(f"\n   Total Texts: {metrics['total_texts']}")
    print(f"   Throughput: {metrics['throughput_texts_per_sec']:.2f} texts/sec")
    print(f"   Avg Latency: {metrics['avg_latency_ms']:.2f} ms")

    if metrics.get('latency'):
        print(f"\n   Latency Percentiles:")
        print(f"     p50: {metrics['latency']['p50']:.2f} ms")
        print(f"     p95: {metrics['latency']['p95']:.2f} ms")
        print(f"     p99: {metrics['latency']['p99']:.2f} ms")

    if metrics.get('cost'):
        print(f"\n   Cost Metrics:")
        print(f"     Total Cost: ${metrics['cost']['total_cost_usd']:.4f}")
        print(f"     Total Tokens: {metrics['cost']['total_tokens']:,}")
        print(f"     Avg Cost/Text: ${metrics['cost']['avg_cost_per_text']:.6f}")

    # Example 4: Full dashboard
    print("\n4. Full Dashboard:")
    monitored.print_dashboard()

    print("\n" + "=" * 60)
    print("Demo complete!")
    print("\nKey Metrics to Track:")
    print("  1. Latency (p95, p99 - not just average)")
    print("  2. Throughput (texts/sec)")
    print("  3. Cache hit rate (target: >60%)")
    print("  4. Cost (total USD and per-text)")
    print("  5. Error rate (target: <1%)")


if __name__ == "__main__":
    main()