ccproxy-api 0.1.6__py3-none-any.whl → 0.2.0__py3-none-any.whl

ccproxy/plugins/max_tokens/service.py

@@ -0,0 +1,271 @@
+"""Service for managing token limits and max_tokens modifications."""
+
+import json
+from pathlib import Path
+from typing import Any
+
+from ccproxy.core.logging import get_plugin_logger
+
+from .config import MaxTokensConfig
+from .models import MaxTokensModification, ModelTokenLimits, TokenLimitsData
+
+
+logger = get_plugin_logger(__name__)
+
+
+class TokenLimitsService:
+    """Service for managing model token limits and max_tokens modifications."""
+
+    def __init__(self, config: MaxTokensConfig):
+        """Initialize token limits service."""
+        self.config = config
+        self.token_limits_data = TokenLimitsData()
+        self._pricing_cache_path = (
+            Path.home() / ".cache" / "ccproxy" / "model_pricing.json"
+        )
+
+        if self.config.prioritize_local_file:
+            # Load local file first (takes precedence)
+            self._load_limits_from_local_file()
+            self._load_limits_from_pricing_cache()
+        else:
+            # Load pricing cache first, local file as fallback
+            self._load_limits_from_pricing_cache()
+            self._load_limits_from_local_file()
+
+    def _load_limits_from_pricing_cache(self) -> None:
+        """Load token limits from pricing plugin cache."""
+        logger.debug(
+            "loading_token_limits_from_pricing_cache",
+            cache_path=str(self._pricing_cache_path),
+        )
+
+        if not self._pricing_cache_path.exists():
+            logger.warning(
+                "pricing_cache_not_found_plugin_will_not_modify_requests",
+                cache_path=str(self._pricing_cache_path),
+                message="max_tokens plugin requires pricing cache to operate",
+            )
+            return
+
+        try:
+            with self._pricing_cache_path.open("r", encoding="utf-8") as f:
+                pricing_data = json.load(f)
+
+            loaded_count = 0
+            for model_name, model_data in pricing_data.items():
+                # Skip non-dict entries (like headers or metadata)
+                if not isinstance(model_data, dict):
+                    continue
+
+                # Skip image generation models and other non-text models
+                if model_data.get("mode") == "image_generation":
+                    continue
+
+                # Extract max_output_tokens (prefer this over max_tokens)
+                max_output = model_data.get("max_output_tokens") or model_data.get(
+                    "max_tokens"
+                )
+                max_input = model_data.get("max_input_tokens")
+
+                # Skip if values are not integers (e.g., documentation strings)
+                if not isinstance(max_output, int):
+                    continue
+
+                if max_output:
+                    self.token_limits_data.models[model_name] = ModelTokenLimits(
+                        max_output_tokens=max_output,
+                        max_input_tokens=max_input
+                        if isinstance(max_input, int)
+                        else None,
+                    )
+                    loaded_count += 1
+
+            logger.debug(
+                "token_limits_loaded_from_pricing_cache",
+                model_count=loaded_count,
+                cache_path=str(self._pricing_cache_path),
+            )
+
+        except Exception as e:
+            logger.error(
+                "failed_to_load_pricing_cache_plugin_will_not_modify_requests",
+                cache_path=str(self._pricing_cache_path),
+                error=str(e),
+                exc_info=e,
+            )
+
+    def _load_limits_from_local_file(self) -> None:
+        """Load token limits from local token_limits.json file."""
+        local_file_path = Path(self.config.default_token_limits_file)
+
+        logger.debug(
+            "loading_token_limits_from_local_file",
+            file_path=str(local_file_path),
+        )
+
+        if not local_file_path.exists():
+            logger.debug(
+                "local_token_limits_file_not_found",
+                file_path=str(local_file_path),
+            )
+            return
+
+        try:
+            with local_file_path.open("r", encoding="utf-8") as f:
+                local_data = json.load(f)
+
+            # Handle flat structure like pricing cache (no nested "models" object)
+            models_data = {}
+            if "models" in local_data:
+                # Old format with nested models
+                models_data = local_data.get("models", {})
+                if not isinstance(models_data, dict):
+                    logger.warning(
+                        "invalid_local_token_limits_format",
+                        file_path=str(local_file_path),
+                        reason="models section is not a dictionary",
+                    )
+                    return
+            else:
+                # New flat format like pricing cache
+                models_data = {
+                    k: v
+                    for k, v in local_data.items()
+                    if not k.startswith("_") and isinstance(v, dict)
+                }
+
+            loaded_count = 0
+            for model_name, model_limits in models_data.items():
+                if not isinstance(model_limits, dict):
+                    continue
+
+                max_output = model_limits.get("max_output_tokens")
+                max_input = model_limits.get("max_input_tokens")
+
+                if isinstance(max_output, int) and max_output > 0:
+                    if self.config.prioritize_local_file:
+                        # Local file values take precedence over pricing cache
+                        self.token_limits_data.models[model_name] = ModelTokenLimits(
+                            max_output_tokens=max_output,
+                            max_input_tokens=max_input
+                            if isinstance(max_input, int) and max_input > 0
+                            else None,
+                        )
+                        loaded_count += 1
+                    else:
+                        # Local file is fallback - only add if model doesn't exist
+                        if model_name not in self.token_limits_data.models:
+                            self.token_limits_data.models[model_name] = (
+                                ModelTokenLimits(
+                                    max_output_tokens=max_output,
+                                    max_input_tokens=max_input
+                                    if isinstance(max_input, int) and max_input > 0
+                                    else None,
+                                )
+                            )
+                            loaded_count += 1
+
+            logger.debug(
+                "token_limits_loaded_from_local_file",
+                file_path=str(local_file_path),
+                model_count=loaded_count,
+            )
+
+        except Exception as e:
+            logger.error(
+                "failed_to_load_local_token_limits_file",
+                file_path=str(local_file_path),
+                error=str(e),
+                exc_info=e,
+            )
+
+    def get_max_output_tokens(self, model_name: str) -> int | None:
+        """Get maximum output tokens for a model."""
+        return self.token_limits_data.get_max_output_tokens(model_name)
+
+    def should_modify_max_tokens(
+        self, request_data: dict[str, Any], model: str
+    ) -> tuple[bool, str]:
+        """Determine if max_tokens should be modified for the request."""
+        current_max_tokens = request_data.get("max_tokens")
+
+        # Enforce mode: always modify to set max_tokens to model limit
+        if self.config.enforce_mode:
+            return True, "enforced"
+
+        # Case 1: No max_tokens provided
+        if current_max_tokens is None:
+            return True, "missing"
+
+        # Case 2: Invalid max_tokens (not a positive integer)
+        if not isinstance(current_max_tokens, int) or current_max_tokens <= 0:
+            return True, "invalid"
+
+        # Case 3: Max tokens exceeds model limit
+        model_limit = self.get_max_output_tokens(model)
+        if model_limit and current_max_tokens > model_limit:
+            return True, "exceeded"
+
+        # No modification needed
+        return False, "none"
+
+    def modify_max_tokens(
+        self, request_data: dict[str, Any], model: str, provider: str | None = None
+    ) -> tuple[dict[str, Any], MaxTokensModification | None]:
+        """Modify max_tokens in request data if needed."""
+        should_modify, reason_type = self.should_modify_max_tokens(request_data, model)
+
+        if not should_modify:
+            return request_data, None
+
+        original_max_tokens = request_data.get("max_tokens")
+
+        # Determine the appropriate max_tokens value
+        model_limit = self.get_max_output_tokens(model)
+
+        if model_limit:
+            new_max_tokens = model_limit
+        else:
+            # Use fallback when model limit is unknown
+            new_max_tokens = self.config.fallback_max_tokens
+            logger.debug(
+                "using_fallback_max_tokens",
+                model=model,
+                fallback=self.config.fallback_max_tokens,
+            )
+
+        # Create modification info
+        modification = MaxTokensModification(
+            original_max_tokens=original_max_tokens,
+            new_max_tokens=new_max_tokens,
+            model=model,
+            reason=self.config.get_modification_reason(reason_type),
+        )
+
+        # Create modified request data
+        modified_data = request_data.copy()
+        modified_data["max_tokens"] = new_max_tokens
+
+        if self.config.log_modifications:
+            logger.info(
+                "max_tokens_modified",
+                model=model,
+                provider=provider,
+                original=original_max_tokens,
+                new=new_max_tokens,
+                reason=modification.reason,
+            )
+
+        return modified_data, modification
+
+    def initialize(self) -> None:
+        """Initialize the service."""
+        logger.debug(
+            "token_limits_service_initialized",
+            models_count=len(self.token_limits_data.models),
+            pricing_cache=str(self._pricing_cache_path),
+            fallback=self.config.fallback_max_tokens,
+            enforce_mode=self.config.enforce_mode,
+            prioritize_local_file=self.config.prioritize_local_file,
+        )

ccproxy/plugins/max_tokens/token_limits.json

@@ -0,0 +1,54 @@
+{
+  "claude-opus-4-1-20250805": {
+    "max_output_tokens": 32000,
+    "max_input_tokens": 200000
+  },
+  "claude-opus-4-20250514": {
+    "max_output_tokens": 32000,
+    "max_input_tokens": 200000
+  },
+  "claude-sonnet-4-20250514": {
+    "max_output_tokens": 64000,
+    "max_input_tokens": 1000000
+  },
+  "claude-3-7-sonnet-20250219": {
+    "max_output_tokens": 8192,
+    "max_input_tokens": 200000
+  },
+  "claude-3-5-sonnet-20241022": {
+    "max_output_tokens": 8192,
+    "max_input_tokens": 200000
+  },
+  "claude-3-5-haiku-20241022": {
+    "max_output_tokens": 8192,
+    "max_input_tokens": 200000
+  },
+  "claude-3-opus-20240229": {
+    "max_output_tokens": 4096,
+    "max_input_tokens": 200000
+  },
+  "claude-3-sonnet-20240229": {
+    "max_output_tokens": 4096,
+    "max_input_tokens": 200000
+  },
+  "claude-3-haiku-20240307": {
+    "max_output_tokens": 4096,
+    "max_input_tokens": 200000
+  },
+  "gpt-5": {
+    "max_output_tokens": 128000,
+    "max_input_tokens": 272000
+  },
+  "gpt-5-codex": {
+    "max_output_tokens": 128000,
+    "max_input_tokens": 272000
+  },
+  "_metadata": {
+    "source": "generated from ~/.cache/ccproxy/model_pricing.json",
+    "claude_models_count": 9,
+    "codex_models_count": 2,
+    "total_models": 11,
+    "generated_for": "max_tokens plugin enforce mode support",
+    "note": "Flat structure format, uses simple model names for compatibility with request handling"
+  }
+}

ccproxy/plugins/metrics/README.md

@@ -0,0 +1,35 @@
+# Metrics Plugin
+
+Collects Prometheus-style metrics for CCProxy events and optionally pushes them.
+
+## Highlights
+- Registers a metrics hook to observe request, token, error, and pool events
+- Exposes a `/metrics` endpoint for scraping when enabled
+- Can schedule a Pushgateway task for off-box metrics delivery
+
+## Configuration
+- `MetricsConfig` toggles collection, namespace, endpoint, and pushgateway options
+- Scheduler integration is automatic when push mode is enabled
+- Generate defaults with `python3 scripts/generate_config_from_model.py \
+  --format toml --plugin metrics --config-class MetricsConfig`
+
+```toml
+[plugins.metrics]
+# enabled = true
+# namespace = "ccproxy"
+# metrics_endpoint_enabled = true
+# pushgateway_enabled = false
+# pushgateway_job = "ccproxy"
+# pushgateway_push_interval = 60
+# collect_request_metrics = true
+# collect_token_metrics = true
+# collect_cost_metrics = true
+# collect_error_metrics = true
+# collect_pool_metrics = true
+# histogram_buckets = [0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 25.0]
+```
+
+## Related Components
+- `hook.py`: implements the metrics collector and event handling
+- `routes.py`: builds the FastAPI router for Prometheus scrape format
+- `tasks.py`: Pushgateway task invoked by the scheduler

ccproxy/plugins/metrics/__init__.py

@@ -0,0 +1,10 @@
+"""Metrics plugin for CCProxy.
+
+This plugin provides Prometheus metrics collection and export functionality
+using the hook system for event-driven metric updates.
+"""
+
+from .plugin import factory
+
+
+__all__ = ["factory"]

ccproxy/{observability/metrics.py → plugins/metrics/collector.py}

@@ -1,5 +1,5 @@
 """
-Prometheus metrics for operational monitoring.
+Prometheus metrics collector for the metrics plugin.
 
 This module provides direct prometheus_client integration for fast operational metrics
 like request counts, response times, and resource usage. These metrics are optimized
@@ -10,7 +10,7 @@ Key features:
 - Minimal overhead for high-frequency operations
 - Standard Prometheus metric types (Counter, Histogram, Gauge)
 - Automatic label management and validation
-- Pushgateway integration for batch metric pushing
+- Integration with hook events for metric updates
 """
 
 from __future__ import annotations
@@ -84,10 +84,10 @@ except ImportError:
     CollectorRegistry = _DummyCollectorRegistry  # type: ignore[misc,assignment]
 
 
-from structlog import get_logger
+from ccproxy.core.logging import get_plugin_logger
 
 
-logger = get_logger(__name__)
+logger = get_plugin_logger(__name__)
 
 
 class PrometheusMetrics:
@@ -102,7 +102,7 @@ class PrometheusMetrics:
         self,
         namespace: str = "ccproxy",
         registry: CollectorRegistry | None = None,
-        pushgateway_client: Any | None = None,
+        histogram_buckets: list[float] | None = None,
     ):
         """
         Initialize Prometheus metrics.
@@ -110,7 +110,7 @@ class PrometheusMetrics:
         Args:
             namespace: Metric name prefix
             registry: Custom Prometheus registry (uses default if None)
-            pushgateway_client: Optional pushgateway client for dependency injection
+            histogram_buckets: Custom histogram bucket boundaries
         """
         if not PROMETHEUS_AVAILABLE:
             logger.warning(
@@ -127,13 +127,21 @@ class PrometheusMetrics:
         else:
             self.registry = registry
         self._enabled = PROMETHEUS_AVAILABLE
-        self._pushgateway_client = pushgateway_client
+        self._histogram_buckets = histogram_buckets or [
+            0.01,
+            0.05,
+            0.1,
+            0.25,
+            0.5,
+            1.0,
+            2.5,
+            5.0,
+            10.0,
+            25.0,
+        ]
 
         if self._enabled:
             self._init_metrics()
-            # Initialize pushgateway client if not provided via DI
-            if self._pushgateway_client is None:
-                self._init_pushgateway()
 
     def _init_metrics(self) -> None:
         """Initialize all Prometheus metric objects."""
@@ -149,7 +157,7 @@ class PrometheusMetrics:
             f"{self.namespace}_response_duration_seconds",
             "Response time in seconds",
             labelnames=["model", "endpoint", "service_type"],
-            buckets=[0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 25.0],
+            buckets=self._histogram_buckets,
             registry=self.registry,
         )
 
@@ -263,7 +271,7 @@ class PrometheusMetrics:
 
         # Set initial system info
         try:
-            from ccproxy import __version__
+            from ccproxy.core import __version__
 
             version = __version__
         except ImportError:
@@ -279,28 +287,6 @@ class PrometheusMetrics:
         # Set service as up
         self.up.labels(job="ccproxy").set(1)
 
-    def _init_pushgateway(self) -> None:
-        """Initialize Pushgateway client if configured (fallback for non-DI usage)."""
-        try:
-            # Import here to avoid circular imports
-            from ccproxy.config.settings import get_settings
-
-            from .pushgateway import PushgatewayClient
-
-            settings = get_settings()
-
-            self._pushgateway_client = PushgatewayClient(settings.observability)
-
-            if self._pushgateway_client.is_enabled():
-                logger.info(
-                    "pushgateway_initialized: url=%s job=%s",
-                    settings.observability.pushgateway_url,
-                    settings.observability.pushgateway_job,
-                )
-        except Exception as e:
-            logger.warning("pushgateway_init_failed: error=%s", str(e))
-            self._pushgateway_client = None
-
     def record_request(
         self,
         method: str,
@@ -473,57 +459,6 @@ class PrometheusMetrics:
         """Check if metrics collection is enabled."""
         return self._enabled
 
-    def push_to_gateway(self, method: str = "push") -> bool:
-        """
-        Push current metrics to Pushgateway using official prometheus_client methods.
-
-        Args:
-            method: Push method - "push" (replace), "pushadd" (add), or "delete"
-
-        Returns:
-            True if push succeeded, False otherwise
-        """
-
-        if not self._enabled or not self._pushgateway_client:
-            return False
-
-        result = self._pushgateway_client.push_metrics(self.registry, method)
-        return bool(result)
-
-    def push_add_to_gateway(self) -> bool:
-        """
-        Add current metrics to existing job/instance in Pushgateway (pushadd operation).
-
-        This is useful when you want to add metrics without replacing existing ones.
-
-        Returns:
-            True if push succeeded, False otherwise
-        """
-        return self.push_to_gateway(method="pushadd")
-
-    def delete_from_gateway(self) -> bool:
-        """
-        Delete all metrics for the configured job from Pushgateway.
-
-        This removes all metrics associated with the job, useful for cleanup.
-
-        Returns:
-            True if delete succeeded, False otherwise
-        """
-
-        if not self._enabled or not self._pushgateway_client:
-            return False
-
-        result = self._pushgateway_client.delete_metrics()
-        return bool(result)
-
-    def is_pushgateway_enabled(self) -> bool:
-        """Check if Pushgateway client is enabled and configured."""
-        return (
-            self._pushgateway_client is not None
-            and self._pushgateway_client.is_enabled()
-        )
-
     # Claude SDK Pool metrics methods
 
     def update_pool_gauges(
@@ -618,71 +553,3 @@ class PrometheusMetrics:
             return
 
         self.pool_clients_active.set(count)
-
-
-# Global metrics instance
-_global_metrics: PrometheusMetrics | None = None
-
-
-def get_metrics(
-    namespace: str = "ccproxy",
-    registry: CollectorRegistry | None = None,
-    pushgateway_client: Any | None = None,
-    settings: Any | None = None,
-) -> PrometheusMetrics:
-    """
-    Get or create global metrics instance with dependency injection.
-
-    Args:
-        namespace: Metric namespace prefix
-        registry: Custom Prometheus registry
-        pushgateway_client: Optional pushgateway client for dependency injection
-        settings: Optional settings instance to avoid circular imports
-
-    Returns:
-        PrometheusMetrics instance with full pushgateway support:
-        - push_to_gateway(): Replace all metrics (default)
-        - push_add_to_gateway(): Add metrics to existing job
-        - delete_from_gateway(): Delete all metrics for job
-    """
-    global _global_metrics
-
-    if _global_metrics is None:
-        # Create pushgateway client if not provided via DI
-        if pushgateway_client is None:
-            from .pushgateway import get_pushgateway_client
-
-            pushgateway_client = get_pushgateway_client()
-
-        _global_metrics = PrometheusMetrics(
-            namespace=namespace,
-            registry=registry,
-            pushgateway_client=pushgateway_client,
-        )
-
-    return _global_metrics
-
-
-def reset_metrics() -> None:
-    """Reset global metrics instance (mainly for testing)."""
-    global _global_metrics
-    _global_metrics = None
-
-    # Clear Prometheus registry to avoid duplicate metrics in tests
-    if PROMETHEUS_AVAILABLE:
-        try:
-            from prometheus_client import REGISTRY
-
-            # Clear all collectors from the registry
-            collectors = list(REGISTRY._collector_to_names.keys())
-            for collector in collectors:
-                REGISTRY.unregister(collector)
-        except Exception:
-            # If clearing the registry fails, just continue
-            # This is mainly for testing and shouldn't break functionality
-            pass
-
-    # Also reset pushgateway client
-    from .pushgateway import reset_pushgateway_client
-
-    reset_pushgateway_client()