ccproxy-api 0.1.0__py3-none-any.whl

ccproxy/testing/response_handlers.py

@@ -0,0 +1,161 @@
+"""Response processing utilities for testing."""
+
+import json
+from typing import Any
+
+import httpx
+
+from ccproxy.testing.config import RequestScenario
+
+
+class ResponseHandler:
+    """Handle responses from both Anthropic and OpenAI formats."""
+
+    def process_response(
+        self, response: httpx.Response, scenario: RequestScenario
+    ) -> dict[str, Any]:
+        """Process response based on format and streaming."""
+
+        if scenario.streaming:
+            return self._process_streaming_response(response, scenario)
+        else:
+            return self._process_standard_response(response, scenario)
+
+    def _process_standard_response(
+        self, response: httpx.Response, scenario: RequestScenario
+    ) -> dict[str, Any]:
+        """Process non-streaming response."""
+
+        try:
+            response_data = response.json()
+
+            # Extract metrics based on format
+            if scenario.api_format == "openai":
+                tokens_input = response_data.get("usage", {}).get("prompt_tokens")
+                tokens_output = response_data.get("usage", {}).get("completion_tokens")
+                content = (
+                    response_data.get("choices", [{}])[0]
+                    .get("message", {})
+                    .get("content")
+                )
+            else:  # anthropic
+                usage = response_data.get("usage", {})
+                tokens_input = usage.get("input_tokens")
+                tokens_output = usage.get("output_tokens")
+                content = ""
+                for block in response_data.get("content", []):
+                    if block.get("type") == "text":
+                        content += block.get("text", "")
+
+            return {
+                "status_code": response.status_code,
+                "headers": dict(response.headers),
+                "data": response_data,
+                "tokens_input": tokens_input,
+                "tokens_output": tokens_output,
+                "content_preview": content[:100] if content else None,
+                "format": scenario.api_format,
+            }
+
+        except Exception as e:
+            return {
+                "status_code": response.status_code,
+                "headers": dict(response.headers),
+                "error": f"Failed to parse {scenario.api_format} response: {str(e)}",
+                "raw_text": response.text[:500] if hasattr(response, "text") else "",
+            }
+
+    def _process_streaming_response(
+        self, response: httpx.Response, scenario: RequestScenario
+    ) -> dict[str, Any]:
+        """Process streaming response."""
+
+        chunks = []
+        total_content = ""
+
+        try:
+            for line in response.iter_lines():
+                if line.startswith("data: "):
+                    data_str = line[6:].strip()
+                    if data_str and data_str != "[DONE]":
+                        try:
+                            chunk_data = json.loads(data_str)
+                            chunks.append(chunk_data)
+
+                            # Extract content based on format
+                            if scenario.api_format == "openai":
+                                delta_content = (
+                                    chunk_data.get("choices", [{}])[0]
+                                    .get("delta", {})
+                                    .get("content", "")
+                                )
+                                total_content += delta_content
+                            else:  # anthropic
+                                if chunk_data.get("type") == "content_block_delta":
+                                    delta_text = chunk_data.get("delta", {}).get(
+                                        "text", ""
+                                    )
+                                    total_content += delta_text
+                        except json.JSONDecodeError:
+                            continue
+
+            return {
+                "status_code": response.status_code,
+                "headers": dict(response.headers),
+                "chunks": chunks,
+                "chunk_count": len(chunks),
+                "total_content": total_content,
+                "content_preview": total_content[:100] if total_content else None,
+                "format": scenario.api_format,
+            }
+
+        except Exception as e:
+            return {
+                "status_code": response.status_code,
+                "headers": dict(response.headers),
+                "error": f"Failed to process {scenario.api_format} stream: {str(e)}",
+            }
+
+
+class MetricsExtractor:
+    """Extract metrics from API responses."""
+
+    @staticmethod
+    def extract_token_metrics(
+        response_data: dict[str, Any], api_format: str
+    ) -> dict[str, int | None]:
+        """Extract token usage from response data."""
+        if api_format == "openai":
+            usage = response_data.get("usage", {})
+            return {
+                "input_tokens": usage.get("prompt_tokens"),
+                "output_tokens": usage.get("completion_tokens"),
+                "cache_read_tokens": None,  # OpenAI doesn't expose cache metrics
+                "cache_write_tokens": None,
+            }
+        else:  # anthropic
+            usage = response_data.get("usage", {})
+            return {
+                "input_tokens": usage.get("input_tokens"),
+                "output_tokens": usage.get("output_tokens"),
+                "cache_read_tokens": usage.get("cache_read_input_tokens"),
+                "cache_write_tokens": usage.get("cache_creation_input_tokens"),
+            }
+
+    @staticmethod
+    def extract_content(response_data: dict[str, Any], api_format: str) -> str:
+        """Extract text content from response data."""
+        if api_format == "openai":
+            content = (
+                response_data.get("choices", [{}])[0]
+                .get("message", {})
+                .get("content", "")
+            )
+            return content if isinstance(content, str) else ""
+        else:  # anthropic
+            content = ""
+            for block in response_data.get("content", []):
+                if block.get("type") == "text":
+                    text = block.get("text", "")
+                    content += text if isinstance(text, str) else ""
+            return content

ccproxy/testing/scenarios.py

@@ -0,0 +1,241 @@
+"""Scenario generation and traffic pattern utilities."""
+
+import random
+from datetime import UTC, datetime, timedelta
+from typing import Any, Literal
+
+from ccproxy.testing.config import RequestScenario, ResponseType, TrafficConfig
+
+
+class ScenarioGenerator:
+    """Generate request scenarios based on traffic configuration."""
+
+    def __init__(self, config: TrafficConfig):
+        self.config = config
+
+    def generate_scenarios(self) -> list[RequestScenario]:
+        """Generate request scenarios based on configuration."""
+        total_requests = int(
+            self.config.duration_seconds * self.config.requests_per_second
+        )
+        scenarios = []
+
+        # Calculate timeframe
+        start_time = self.config.start_timestamp or datetime.now(UTC)
+        time_span = self.config.duration_seconds
+
+        for i in range(total_requests):
+            # Determine timing based on pattern
+            time_offset = self._calculate_time_offset(i, total_requests, time_span)
+            request_time = start_time + time_offset
+
+            # Select random parameters
+            model = random.choice(self.config.models)
+            message_type = random.choice(self.config.message_types)
+            streaming = random.random() < self.config.streaming_probability
+
+            # Determine response type
+            response_type = self._determine_response_type()
+
+            # Determine API format based on distribution
+            api_format = self._determine_api_format()
+
+            # Set endpoint path based on format
+            endpoint_path = (
+                "/api/v1/chat/completions"
+                if api_format == "openai"
+                else "/api/v1/messages"
+            )
+
+            # Generate headers with bypass and format-specific headers
+            headers = self._generate_headers(api_format, streaming)
+
+            scenarios.append(
+                RequestScenario(
+                    model=model,
+                    message_type=message_type,
+                    streaming=streaming,
+                    response_type=response_type,
+                    timestamp=request_time,
+                    api_format=api_format,
+                    endpoint_path=endpoint_path,
+                    bypass_upstream=self.config.bypass_mode,
+                    use_real_auth=not self.config.bypass_mode,
+                    headers=headers,
+                    target_url=self.config.target_url,
+                )
+            )
+
+        return scenarios
+
+    def _calculate_time_offset(
+        self, request_index: int, total_requests: int, time_span: int
+    ) -> timedelta:
+        """Calculate time offset for request based on traffic pattern."""
+        if self.config.pattern == "constant":
+            return timedelta(seconds=request_index / self.config.requests_per_second)
+        elif self.config.pattern == "burst":
+            # Front-load requests in bursts
+            burst_size = max(1, int(total_requests * 0.1))
+            if request_index < burst_size:
+                return timedelta(seconds=request_index * 0.1)
+            else:
+                remaining_time = time_span - (burst_size * 0.1)
+                remaining_requests = total_requests - burst_size
+                return timedelta(
+                    seconds=(burst_size * 0.1)
+                    + ((request_index - burst_size) / remaining_requests)
+                    * remaining_time
+                )
+        elif self.config.pattern == "ramping":
+            # Gradually increase request rate
+            normalized_time = request_index / total_requests
+            accelerated_time = normalized_time**2
+            return timedelta(seconds=accelerated_time * time_span)
+        else:  # realistic
+            # Add some randomness to simulate real user behavior
+            base_time = request_index / self.config.requests_per_second
+            jitter = random.uniform(-0.5, 0.5)
+            return timedelta(seconds=max(0, base_time + jitter))
+
+    def _determine_response_type(self) -> ResponseType:
+        """Determine response type based on configuration."""
+        if self.config.response_type == "mixed":
+            rand = random.random()
+            if rand < self.config.error_probability:
+                return "error"
+            elif rand < self.config.error_probability * 1.2:
+                return "unavailable"
+            else:
+                return "success"
+        else:
+            return self.config.response_type
+
+    def _determine_api_format(self) -> Literal["anthropic", "openai"]:
+        """Determine API format based on distribution configuration."""
+        if len(self.config.api_formats) == 1:
+            format_name = self.config.api_formats[0]
+            if format_name == "anthropic":
+                return "anthropic"
+            elif format_name == "openai":
+                return "openai"
+            return "anthropic"  # Default fallback
+
+        # Use weighted random selection based on format_distribution
+        rand = random.random()
+        cumulative = 0.0
+
+        for format_name in self.config.api_formats:
+            weight = self.config.format_distribution.get(format_name, 0.0)
+            cumulative += weight
+            if rand <= cumulative:
+                if format_name == "anthropic":
+                    return "anthropic"
+                elif format_name == "openai":
+                    return "openai"
+
+        # Fallback to first format if distribution doesn't add up
+        format_name = self.config.api_formats[0]
+        if format_name == "anthropic":
+            return "anthropic"
+        elif format_name == "openai":
+            return "openai"
+        return "anthropic"  # Default fallback
+
+    def _generate_headers(self, api_format: str, streaming: bool) -> dict[str, str]:
+        """Generate headers with bypass and format-specific headers."""
+        headers = {}
+
+        # Add bypass header if in bypass mode
+        if self.config.bypass_mode:
+            headers["X-CCProxy-Bypass-Upstream"] = "true"
+
+        # Add real API authentication if not in bypass mode
+        if not self.config.bypass_mode and self.config.real_api_keys:
+            if api_format == "openai" and "openai" in self.config.real_api_keys:
+                headers["Authorization"] = (
+                    f"Bearer {self.config.real_api_keys['openai']}"
+                )
+            elif api_format == "anthropic" and "anthropic" in self.config.real_api_keys:
+                headers["Authorization"] = (
+                    f"Bearer {self.config.real_api_keys['anthropic']}"
+                )
+
+        # Format-specific headers
+        if api_format == "openai":
+            headers["Content-Type"] = "application/json"
+            headers["Accept"] = "application/json"
+        else:  # anthropic
+            headers["Content-Type"] = "application/json"
+            headers["Accept"] = "application/json"
+            headers["anthropic-version"] = "2023-06-01"
+
+        # Streaming-specific headers
+        if streaming:
+            headers["Accept"] = "text/event-stream"
+            headers["Cache-Control"] = "no-cache"
+
+        return headers
+
+
+class TrafficPatternAnalyzer:
+    """Analyze and validate traffic patterns."""
+
+    @staticmethod
+    def analyze_distribution(scenarios: list[RequestScenario]) -> dict[str, Any]:
+        """Analyze the distribution of scenarios."""
+        analysis = {
+            "total_scenarios": len(scenarios),
+            "api_format_distribution": {},
+            "model_distribution": {},
+            "message_type_distribution": {},
+            "streaming_percentage": 0.0,
+            "time_span_seconds": 0.0,
+        }
+
+        if not scenarios:
+            return analysis
+
+        # Count distributions
+        api_formats: dict[str, int] = {}
+        models: dict[str, int] = {}
+        message_types: dict[str, int] = {}
+        streaming_count = 0
+
+        for scenario in scenarios:
+            # API format distribution
+            api_formats[scenario.api_format] = (
+                api_formats.get(scenario.api_format, 0) + 1
+            )
+
+            # Model distribution
+            models[scenario.model] = models.get(scenario.model, 0) + 1
+
+            # Message type distribution
+            message_types[scenario.message_type] = (
+                message_types.get(scenario.message_type, 0) + 1
+            )
+
+            # Streaming count
+            if scenario.streaming:
+                streaming_count += 1
+
+        # Calculate percentages
+        total = len(scenarios)
+        analysis["api_format_distribution"] = {
+            k: v / total for k, v in api_formats.items()
+        }
+        analysis["model_distribution"] = {k: v / total for k, v in models.items()}
+        analysis["message_type_distribution"] = {
+            k: v / total for k, v in message_types.items()
+        }
+        analysis["streaming_percentage"] = streaming_count / total
+
+        # Calculate time span
+        timestamps = [scenario.timestamp for scenario in scenarios]
+        if timestamps:
+            analysis["time_span_seconds"] = (
+                max(timestamps) - min(timestamps)
+            ).total_seconds()
+
+        return analysis

ccproxy/utils/__init__.py

@@ -0,0 +1,6 @@
+"""Utility modules for shared functionality across the application."""
+
+from .cost_calculator import calculate_cost_breakdown, calculate_token_cost
+
+
+__all__ = ["calculate_token_cost", "calculate_cost_breakdown"]

ccproxy/utils/cost_calculator.py

@@ -0,0 +1,210 @@
+"""Cost calculation utilities for token-based pricing.
+
+This module provides shared cost calculation functionality that can be used
+across different services to ensure consistent pricing calculations.
+"""
+
+import structlog
+
+
+logger = structlog.get_logger(__name__)
+
+
+def calculate_token_cost(
+    tokens_input: int | None,
+    tokens_output: int | None,
+    model: str | None,
+    cache_read_tokens: int | None = None,
+    cache_write_tokens: int | None = None,
+) -> float | None:
+    """Calculate cost in USD for the given token usage including cache tokens.
+
+    This is a shared utility function that provides consistent cost calculation
+    across all services using the pricing data from the pricing system.
+
+    Args:
+        tokens_input: Number of input tokens
+        tokens_output: Number of output tokens
+        model: Model name for pricing lookup
+        cache_read_tokens: Number of cache read tokens
+        cache_write_tokens: Number of cache write tokens
+
+    Returns:
+        Cost in USD or None if calculation not possible
+    """
+    if not model or (
+        not tokens_input
+        and not tokens_output
+        and not cache_read_tokens
+        and not cache_write_tokens
+    ):
+        return None
+
+    try:
+        # Import pricing system components
+        from ccproxy.config.pricing import PricingSettings
+        from ccproxy.pricing.cache import PricingCache
+        from ccproxy.pricing.loader import PricingLoader
+
+        # Get canonical model name
+        canonical_model = PricingLoader.get_canonical_model_name(model)
+
+        # Create pricing components with dependency injection
+        settings = PricingSettings()
+        cache = PricingCache(settings)
+        cached_data = cache.load_cached_data()
+
+        # If cache is expired, try to use stale cache as fallback
+        if not cached_data:
+            try:
+                import json
+
+                if cache.cache_file.exists():
+                    with cache.cache_file.open(encoding="utf-8") as f:
+                        cached_data = json.load(f)
+                    logger.debug(
+                        "cost_calculation_using_stale_cache",
+                        cache_age_hours=cache.get_cache_info().get("age_hours"),
+                    )
+            except (OSError, json.JSONDecodeError):
+                pass
+
+        if not cached_data:
+            logger.debug("cost_calculation_skipped", reason="no_pricing_data")
+            return None
+
+        # Load pricing data
+        pricing_data = PricingLoader.load_pricing_from_data(cached_data, verbose=False)
+        if not pricing_data or canonical_model not in pricing_data:
+            logger.debug(
+                "cost_calculation_skipped",
+                model=canonical_model,
+                reason="model_not_found",
+            )
+            return None
+
+        model_pricing = pricing_data[canonical_model]
+
+        # Calculate cost (pricing is per 1M tokens)
+        input_cost = ((tokens_input or 0) / 1_000_000) * float(model_pricing.input)
+        output_cost = ((tokens_output or 0) / 1_000_000) * float(model_pricing.output)
+        cache_read_cost = ((cache_read_tokens or 0) / 1_000_000) * float(
+            model_pricing.cache_read
+        )
+        cache_write_cost = ((cache_write_tokens or 0) / 1_000_000) * float(
+            model_pricing.cache_write
+        )
+
+        total_cost = input_cost + output_cost + cache_read_cost + cache_write_cost
+
+        logger.debug(
+            "cost_calculated",
+            model=canonical_model,
+            tokens_input=tokens_input,
+            tokens_output=tokens_output,
+            cache_read_tokens=cache_read_tokens,
+            cache_write_tokens=cache_write_tokens,
+            input_cost=input_cost,
+            output_cost=output_cost,
+            cache_read_cost=cache_read_cost,
+            cache_write_cost=cache_write_cost,
+            cost_usd=total_cost,
+        )
+
+        return total_cost
+
+    except Exception as e:
+        logger.debug("cost_calculation_error", error=str(e), model=model)
+        return None
+
+
+def calculate_cost_breakdown(
+    tokens_input: int | None,
+    tokens_output: int | None,
+    model: str | None,
+    cache_read_tokens: int | None = None,
+    cache_write_tokens: int | None = None,
+) -> dict[str, float | str] | None:
+    """Calculate detailed cost breakdown for the given token usage.
+
+    Args:
+        tokens_input: Number of input tokens
+        tokens_output: Number of output tokens
+        model: Model name for pricing lookup
+        cache_read_tokens: Number of cache read tokens
+        cache_write_tokens: Number of cache write tokens
+
+    Returns:
+        Dictionary with cost breakdown or None if calculation not possible
+    """
+    if not model or (
+        not tokens_input
+        and not tokens_output
+        and not cache_read_tokens
+        and not cache_write_tokens
+    ):
+        return None
+
+    try:
+        # Import pricing system components
+        from ccproxy.config.pricing import PricingSettings
+        from ccproxy.pricing.cache import PricingCache
+        from ccproxy.pricing.loader import PricingLoader
+
+        # Get canonical model name
+        canonical_model = PricingLoader.get_canonical_model_name(model)
+
+        # Create pricing components with dependency injection
+        settings = PricingSettings()
+        cache = PricingCache(settings)
+        cached_data = cache.load_cached_data()
+
+        # If cache is expired, try to use stale cache as fallback
+        if not cached_data:
+            try:
+                import json
+
+                if cache.cache_file.exists():
+                    with cache.cache_file.open(encoding="utf-8") as f:
+                        cached_data = json.load(f)
+                    logger.debug(
+                        "cost_breakdown_using_stale_cache",
+                        cache_age_hours=cache.get_cache_info().get("age_hours"),
+                    )
+            except (OSError, json.JSONDecodeError):
+                pass
+
+        if not cached_data:
+            return None
+
+        # Load pricing data
+        pricing_data = PricingLoader.load_pricing_from_data(cached_data, verbose=False)
+        if not pricing_data or canonical_model not in pricing_data:
+            return None
+
+        model_pricing = pricing_data[canonical_model]
+
+        # Calculate individual costs (pricing is per 1M tokens)
+        input_cost = ((tokens_input or 0) / 1_000_000) * float(model_pricing.input)
+        output_cost = ((tokens_output or 0) / 1_000_000) * float(model_pricing.output)
+        cache_read_cost = ((cache_read_tokens or 0) / 1_000_000) * float(
+            model_pricing.cache_read
+        )
+        cache_write_cost = ((cache_write_tokens or 0) / 1_000_000) * float(
+            model_pricing.cache_write
+        )
+
+        total_cost = input_cost + output_cost + cache_read_cost + cache_write_cost
+
+        return {
+            "input_cost": input_cost,
+            "output_cost": output_cost,
+            "cache_read_cost": cache_read_cost,
+            "cache_write_cost": cache_write_cost,
+            "total_cost": total_cost,
+            "model": canonical_model,
+        }
+
+    except Exception as e:
+        logger.debug("cost_breakdown_error", error=str(e), model=model)
+        return None