ccproxy-api 0.1.7__py3-none-any.whl → 0.2.0a4__py3-none-any.whl

ccproxy/plugins/max_tokens/adapter.py

@@ -0,0 +1,235 @@
+"""Max tokens adapter implementation using hook system."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from ccproxy.core.logging import get_plugin_logger
+from ccproxy.core.plugins.hooks import Hook, HookContext, HookEvent
+from ccproxy.core.plugins.hooks.layers import HookLayer
+
+from .config import MaxTokensConfig
+from .service import TokenLimitsService
+
+
+logger = get_plugin_logger(__name__)
+
+
+class MaxTokensHook(Hook):
+    """Hook that enforces `max_tokens` limits before provider dispatch."""
+
+    name = "max_tokens"
+    events = [HookEvent.PROVIDER_REQUEST_PREPARED]
+    priority = HookLayer.PROCESSING  # Run before observation hooks
+
+    def __init__(self, config: MaxTokensConfig, service: TokenLimitsService):
+        """Initialize max tokens hook."""
+        self.config = config
+        self.service = service
+
+    async def __call__(self, context: HookContext) -> None:
+        """Adjust token limits in provider request payload when required."""
+        if not self.config.enabled:
+            return
+
+        if context.event is not HookEvent.PROVIDER_REQUEST_PREPARED:
+            return
+
+        provider = context.provider
+        if provider and not self.config.should_process_provider(provider):
+            return
+
+        payload = context.data.get("body")
+        if not isinstance(payload, dict):
+            return
+
+        model = payload.get("model")
+        if not isinstance(model, str):
+            return
+
+        provider_model = context.metadata.get("provider_model")
+        if not isinstance(provider_model, str):
+            provider_model = model
+
+        alias_map = context.metadata.get("_model_alias_map")
+        client_model = context.metadata.get("client_model")
+        if not client_model and isinstance(alias_map, dict):
+            client_model = alias_map.get(provider_model) or alias_map.get(model)
+
+        modifiers = context.data.setdefault("modifiers", {})
+        modified = False
+
+        # Ensure we work with the latest payload returned from the service
+        modified_payload, modification = self.service.modify_max_tokens(
+            payload, provider_model, provider
+        )
+
+        if modification and modification.was_modified():
+            payload = modified_payload
+            context.data["body"] = payload
+            modifiers["max_tokens"] = {
+                "original": modification.original_max_tokens,
+                "new": modification.new_max_tokens,
+                "reason": modification.reason,
+            }
+            modified = True
+        else:
+            payload = modified_payload
+            context.data["body"] = payload
+
+        current_max_output = payload.get("max_output_tokens")
+        provider_limit = self.service.get_max_output_tokens(provider_model)
+        if provider_limit is None and self.config.fallback_max_tokens:
+            provider_limit = self.config.fallback_max_tokens
+
+        new_max_output: int | None = None
+        output_reason: str | None = None
+
+        if isinstance(current_max_output, int) and current_max_output > 0:
+            original_limit = (
+                self.service.get_max_output_tokens(client_model)
+                if client_model
+                else None
+            )
+
+            if (
+                provider_limit
+                and original_limit
+                and client_model
+                and client_model != provider_model
+                and current_max_output == original_limit
+                and provider_limit != original_limit
+            ):
+                new_max_output = provider_limit
+                output_reason = "max_output_tokens_aligned_with_mapped_model"
+            elif provider_limit and current_max_output > provider_limit:
+                new_max_output = provider_limit
+                output_reason = "max_output_tokens_capped_to_provider_limit"
+
+        if new_max_output is not None and new_max_output != current_max_output:
+            payload["max_output_tokens"] = new_max_output
+            modifiers["max_output_tokens"] = {
+                "original": current_max_output,
+                "new": new_max_output,
+                "reason": output_reason,
+            }
+            modified = True
+
+            if self.config.log_modifications:
+                logger.info(
+                    "max_output_tokens_adjusted",
+                    provider=provider,
+                    provider_model=provider_model,
+                    client_model=client_model,
+                    original=current_max_output,
+                    new=new_max_output,
+                    reason=output_reason,
+                )
+
+        if modified:
+            context.data["body"] = payload
+            context.data["body_kind"] = "json"
+            context.data["body_raw"] = json.dumps(payload).encode("utf-8")
+
+
+class MaxTokensAdapter:
+    """Max tokens adapter using hook-based request interception."""
+
+    def __init__(self, config: MaxTokensConfig):
+        """Initialize max tokens adapter."""
+        self.config = config
+        self.service = TokenLimitsService(config)
+        self.hook = MaxTokensHook(config, self.service)
+        self._initialized = False
+
+    async def initialize(self) -> None:
+        """Initialize the adapter and register hooks."""
+        if self._initialized:
+            return
+
+        if not self.config.enabled:
+            logger.debug("max_tokens_adapter_disabled")
+            return
+
+        # Initialize the service
+        self.service.initialize()
+
+        # Register hook with hook manager
+        try:
+            from ccproxy.services.container import ServiceContainer
+
+            container = ServiceContainer.get_current(strict=False)
+            if container:
+                try:
+                    hook_registry = container.get_hook_registry()
+                    if hook_registry:
+                        hook_registry.register(self.hook)
+                        logger.debug(
+                            "max_tokens_hook_registered",
+                            hook_name=self.hook.name,
+                            events=[event.value for event in self.hook.events],
+                            priority=self.hook.priority,
+                        )
+                    else:
+                        logger.warning("no_hook_registry_available")
+                except Exception as service_error:
+                    logger.warning(
+                        "hook_registry_service_not_available",
+                        error=str(service_error),
+                    )
+            else:
+                logger.warning("no_service_container_available")
+
+        except Exception as e:
+            logger.error(
+                "failed_to_register_max_tokens_hook",
+                error=str(e),
+                exc_info=e,
+            )
+            # Continue without hook registration - plugin will be effectively disabled
+
+        self._initialized = True
+        logger.info("max_tokens_adapter_initialized")
+
+    async def cleanup(self) -> None:
+        """Cleanup resources."""
+        if not self._initialized:
+            return
+
+        try:
+            # Unregister hook
+            from ccproxy.services.container import ServiceContainer
+
+            container = ServiceContainer.get_current(strict=False)
+            if container:
+                try:
+                    hook_registry = container.get_hook_registry()
+                    if hook_registry:
+                        hook_registry.unregister(self.hook)
+                        logger.debug("max_tokens_hook_unregistered")
+                except Exception as service_error:
+                    logger.debug(
+                        "hook_registry_service_not_available_during_cleanup",
+                        error=str(service_error),
+                    )
+
+        except Exception as e:
+            logger.error(
+                "failed_to_unregister_max_tokens_hook",
+                error=str(e),
+                exc_info=e,
+            )
+
+        self._initialized = False
+        logger.debug("max_tokens_adapter_cleanup_completed")
+
+    def get_modification_stats(self) -> dict[str, Any]:
+        """Get statistics about max_tokens modifications."""
+        # This could be enhanced to track modification statistics
+        return {
+            "adapter_initialized": self._initialized,
+            "config_enabled": self.config.enabled,
+            "target_providers": self.config.target_providers,
+            "fallback_max_tokens": self.config.fallback_max_tokens,
+        }

ccproxy/plugins/max_tokens/config.py

@@ -0,0 +1,86 @@
+"""Configuration for max_tokens plugin."""
+
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class MaxTokensConfig(BaseModel):
+    """Configuration for the max_tokens plugin."""
+
+    enabled: bool = Field(
+        default=True, description="Whether the max_tokens plugin is enabled"
+    )
+    default_token_limits_file: str = Field(
+        default=str(Path(__file__).parent / "token_limits.json"),
+        description="Path to JSON file containing default token limits",
+    )
+    fallback_max_tokens: int = Field(
+        default=4096,
+        ge=1,
+        description="Fallback max_tokens when model limits are unknown",
+    )
+    apply_to_all_providers: bool = Field(
+        default=True,
+        description="Whether to apply to all providers or only specific ones",
+    )
+    target_providers: list[str] = Field(
+        default_factory=lambda: ["claude_api", "claude_sdk", "codex", "copilot"],
+        description="List of providers to apply max_tokens modifications to",
+    )
+    require_pricing_data: bool = Field(
+        default=False,
+        description=(
+            "If True, only modify requests when pricing data is available. "
+            "If False, use fallback limits when pricing data is not available."
+        ),
+    )
+    log_modifications: bool = Field(
+        default=True, description="Whether to log max_tokens modifications"
+    )
+    enforce_mode: bool = Field(
+        default=False,
+        description=(
+            "When enabled, always set max_tokens to the model's maximum limit, "
+            "ignoring the request's current max_tokens value"
+        ),
+    )
+    prioritize_local_file: bool = Field(
+        default=False,
+        description=(
+            "When enabled, local token_limits.json values take precedence over "
+            "pricing cache values. When disabled, local file is only used as fallback "
+            "when pricing cache is unavailable or model is not found in cache."
+        ),
+    )
+    modification_reasons: dict[str, str] = Field(
+        default_factory=lambda: {
+            "missing": "max_tokens was missing from request",
+            "invalid": "max_tokens was invalid or too high",
+            "exceeded": "max_tokens exceeded model limit",
+            "enforced": "max_tokens enforced to model limit (enforce mode)",
+        },
+        description="Reason templates for modifications",
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_config(cls, data: dict[str, Any]) -> dict[str, Any]:
+        """Validate configuration values."""
+        # Ensure target_providers is a list
+        if isinstance(data.get("target_providers"), str):
+            data["target_providers"] = [data["target_providers"]]
+        return data
+
+    def should_process_provider(self, provider: str) -> bool:
+        """Check if plugin should process requests for given provider."""
+        if self.apply_to_all_providers:
+            return True
+        return provider in self.target_providers
+
+    def get_modification_reason(self, reason_type: str) -> str:
+        """Get modification reason text for given reason type."""
+        return self.modification_reasons.get(
+            reason_type, f"Unknown reason: {reason_type}"
+        )

ccproxy/plugins/max_tokens/models.py

@@ -0,0 +1,53 @@
+"""Pydantic models for max_tokens plugin."""
+
+from pydantic import BaseModel, Field
+
+
+class ModelTokenLimits(BaseModel):
+    """Token limits for a specific model."""
+
+    max_output_tokens: int = Field(
+        ..., ge=1, description="Maximum output tokens for the model"
+    )
+    max_input_tokens: int | None = Field(
+        default=None, ge=1, description="Maximum input tokens for the model"
+    )
+
+
+class TokenLimitsData(BaseModel):
+    """Complete token limits data for all models."""
+
+    models: dict[str, ModelTokenLimits] = Field(
+        default_factory=dict, description="Model name to token limits mapping"
+    )
+
+    def get_max_output_tokens(self, model_name: str) -> int | None:
+        """Get maximum output tokens for a model."""
+        model_limits = self.models.get(model_name)
+        return model_limits.max_output_tokens if model_limits else None
+
+    def get_max_input_tokens(self, model_name: str) -> int | None:
+        """Get maximum input tokens for a model."""
+        model_limits = self.models.get(model_name)
+        return model_limits.max_input_tokens if model_limits else None
+
+    def has_model(self, model_name: str) -> bool:
+        """Check if model limits exist for the given model."""
+        return model_name in self.models
+
+
+class MaxTokensModification(BaseModel):
+    """Information about max_tokens modification made by the plugin."""
+
+    original_max_tokens: int | None = Field(
+        description="Original max_tokens value from request"
+    )
+    new_max_tokens: int | None = Field(
+        description="New max_tokens value after modification"
+    )
+    model: str = Field(description="Model name")
+    reason: str = Field(description="Reason for modification")
+
+    def was_modified(self) -> bool:
+        """Check if max_tokens was modified."""
+        return self.original_max_tokens != self.new_max_tokens

ccproxy/plugins/max_tokens/plugin.py

@@ -0,0 +1,200 @@
+"""Max tokens plugin implementation."""
+
+from typing import Any
+
+from ccproxy.core.logging import get_plugin_logger
+from ccproxy.core.plugins import (
+    PluginManifest,
+    SystemPluginFactory,
+    SystemPluginRuntime,
+)
+from ccproxy.core.plugins.hooks import HookRegistry
+
+from .adapter import MaxTokensHook
+from .config import MaxTokensConfig
+from .service import TokenLimitsService
+
+
+logger = get_plugin_logger(__name__)
+
+
+class MaxTokensRuntime(SystemPluginRuntime):
+    """Runtime for max_tokens plugin."""
+
+    def __init__(self, manifest: PluginManifest):
+        """Initialize runtime."""
+        super().__init__(manifest)
+        self.config: MaxTokensConfig | None = None
+        self.service: TokenLimitsService | None = None
+        self.hook: MaxTokensHook | None = None
+        self.hook_registered = False
+
+    async def _on_initialize(self) -> None:
+        """Initialize the max tokens plugin."""
+        if not self.context:
+            raise RuntimeError("Context not set")
+
+        # Get configuration
+        config = self.context.get("config")
+        if not isinstance(config, MaxTokensConfig):
+            logger.debug("plugin_no_config_using_defaults", category="plugin")
+            # Use default config if none provided
+            self.config = MaxTokensConfig()
+        else:
+            self.config = config
+
+        logger.debug("initializing_max_tokens_plugin", enabled=self.config.enabled)
+
+        if not self.config.enabled:
+            logger.debug("max_tokens_plugin_disabled")
+            return
+
+        # Create and initialize service
+        self.service = TokenLimitsService(self.config)
+        self.service.initialize()
+
+        # Create hook instance
+        self.hook = MaxTokensHook(self.config, self.service)
+
+        # Attempt to register hook with available registry sources
+        hook_registry: HookRegistry | None = None
+
+        # 1. Directly from context (preferred when ServiceContainer wires it)
+        hook_registry = self.context.get("hook_registry")
+
+        # 2. Fallback to service container if available
+        if not hook_registry:
+            container = self.context.get("service_container")
+            if container:
+                try:
+                    hook_registry = container.get_hook_registry()
+                except Exception as container_error:  # pragma: no cover - defensive
+                    logger.debug(
+                        "max_tokens_hook_registry_from_container_failed",
+                        error=str(container_error),
+                    )
+
+        # 3. Fallback to app.state when running inside FastAPI app
+        if not hook_registry:
+            app = self.context.get("app")
+            if app and hasattr(app.state, "hook_registry"):
+                hook_registry = app.state.hook_registry
+
+        if hook_registry and isinstance(hook_registry, HookRegistry):
+            hook_registry.register(self.hook)
+            self.hook_registered = True
+            logger.debug(
+                "max_tokens_hook_registered",
+                providers="*"
+                if self.config.apply_to_all_providers
+                else self.config.target_providers,
+                fallback=self.config.fallback_max_tokens,
+            )
+        else:
+            logger.warning(
+                "max_tokens_hook_registry_unavailable",
+                message="max_tokens adjustments disabled",
+            )
+
+        logger.info(
+            "max_tokens_plugin_initialized",
+            target_providers=self.config.target_providers,
+            fallback_max_tokens=self.config.fallback_max_tokens,
+        )
+
+    async def _on_shutdown(self) -> None:
+        """Shutdown the plugin and cleanup resources."""
+        logger.debug("shutting_down_max_tokens_plugin")
+
+        # Unregister hook if we registered one
+        if self.hook:
+            hook_registry: HookRegistry | None = None
+            if self.context:
+                hook_registry = self.context.get("hook_registry")
+                if not hook_registry:
+                    container = self.context.get("service_container")
+                    if container:
+                        try:
+                            hook_registry = container.get_hook_registry()
+                        except Exception as container_error:  # pragma: no cover
+                            logger.debug(
+                                "max_tokens_hook_registry_from_container_failed_shutdown",
+                                error=str(container_error),
+                            )
+                if not hook_registry:
+                    app = self.context.get("app")
+                    if app and hasattr(app.state, "hook_registry"):
+                        hook_registry = app.state.hook_registry
+
+            if hook_registry and isinstance(hook_registry, HookRegistry):
+                hook_registry.unregister(self.hook)
+                self.hook_registered = False
+                logger.debug("max_tokens_hook_unregistered")
+
+            self.hook = None
+
+        if self.service:
+            self.service = None
+
+        logger.debug("max_tokens_plugin_shutdown_complete")
+
+    async def _get_health_details(self) -> dict[str, Any]:
+        """Get health check details."""
+        try:
+            base_health = {
+                "type": "system",
+                "initialized": self.initialized,
+                "enabled": self.config.enabled if self.config else False,
+                "hook_registered": self.hook_registered,
+            }
+
+            if not self.config or not self.config.enabled:
+                return base_health
+
+            # Add service health info
+            health_details = base_health.copy()
+
+            if self.service:
+                health_details["models_count"] = len(
+                    self.service.token_limits_data.models
+                )
+                health_details["fallback_max_tokens"] = self.config.fallback_max_tokens
+
+            return health_details
+
+        except Exception as e:
+            logger.error("health_check_failed", error=str(e))
+            return {
+                "type": "system",
+                "initialized": self.initialized,
+                "enabled": self.config.enabled if self.config else False,
+                "error": str(e),
+                "hook_registered": self.hook_registered,
+            }
+
+
+class MaxTokensFactory(SystemPluginFactory):
+    """Factory for max_tokens plugin."""
+
+    def __init__(self) -> None:
+        """Initialize factory with manifest."""
+        # Create manifest - max_tokens logic is now integrated into HTTP adapter
+        manifest = PluginManifest(
+            name="max_tokens",
+            version="0.1.0",
+            description="Automatically sets max_tokens based on model limits when missing or invalid",
+            is_provider=False,
+            config_class=MaxTokensConfig,
+            provides=["max_tokens"],  # This plugin provides the max_tokens service
+        )
+
+        # Initialize with manifest
+        super().__init__(manifest)
+
+    def create_runtime(self) -> MaxTokensRuntime:
+        """Create runtime instance."""
+        return MaxTokensRuntime(self.manifest)
+
+
+# Export the factory instance
+factory = MaxTokensFactory()