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

ccproxy/plugins/pricing/plugin.py

@@ -0,0 +1,169 @@
+"""Pricing plugin implementation."""
+
+from typing import Any
+
+from ccproxy.core.logging import get_plugin_logger
+from ccproxy.core.plugins import (
+    PluginManifest,
+    SystemPluginFactory,
+    SystemPluginRuntime,
+)
+
+from .config import PricingConfig
+from .service import PricingService
+from .tasks import PricingCacheUpdateTask
+
+
+logger = get_plugin_logger()
+
+
+class PricingRuntime(SystemPluginRuntime):
+    """Runtime for pricing plugin."""
+
+    def __init__(self, manifest: PluginManifest):
+        """Initialize runtime."""
+        super().__init__(manifest)
+        self.config: PricingConfig | None = None
+        self.service: PricingService | None = None
+        self.update_task: PricingCacheUpdateTask | None = None
+
+    async def _on_initialize(self) -> None:
+        """Initialize the pricing plugin."""
+        if not self.context:
+            raise RuntimeError("Context not set")
+
+        # Get configuration
+        config = self.context.get("config")
+        if not isinstance(config, PricingConfig):
+            logger.debug("plugin_no_config_using_defaults", category="plugin")
+            # Use default config if none provided
+            self.config = PricingConfig()
+        else:
+            self.config = config
+
+        logger.debug("initializing_pricing_plugin", enabled=self.config.enabled)
+
+        # Create pricing service
+        self.service = PricingService(self.config)
+
+        if self.config.enabled:
+            # Initialize the service
+            await self.service.initialize()
+
+            # Register service with plugin registry
+            plugin_registry = self.context.get("plugin_registry")
+            if plugin_registry:
+                plugin_registry.register_service(
+                    "pricing", self.service, self.manifest.name
+                )
+                logger.debug("pricing_service_registered")
+
+            # Create and start pricing update task
+            interval_seconds = self.config.update_interval_hours * 3600
+            self.update_task = PricingCacheUpdateTask(
+                name="pricing_cache_update",
+                interval_seconds=interval_seconds,
+                pricing_service=self.service,
+                enabled=self.config.auto_update,
+                force_refresh_on_startup=self.config.force_refresh_on_startup,
+            )
+
+            await self.update_task.start()
+            logger.debug(
+                "pricing_plugin_initialized",
+                update_interval_hours=self.config.update_interval_hours,
+                auto_update=self.config.auto_update,
+                force_refresh_on_startup=self.config.force_refresh_on_startup,
+            )
+        else:
+            logger.debug("pricing_plugin_disabled")
+
+    async def _on_shutdown(self) -> None:
+        """Shutdown the plugin and cleanup resources."""
+        logger.debug("shutting_down_pricing_plugin")
+
+        # Stop the update task
+        if self.update_task:
+            await self.update_task.stop()
+
+        logger.debug("pricing_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,
+            }
+
+            if not self.config or not self.config.enabled:
+                return base_health
+
+            # Add service-specific health info
+            health_details = base_health.copy()
+
+            if self.service:
+                cache_info = self.service.get_cache_info()
+                health_details.update(
+                    {
+                        "cache_valid": cache_info.get("valid", False),
+                        "cache_age_hours": cache_info.get("age_hours"),
+                        "cache_exists": cache_info.get("exists", False),
+                    }
+                )
+
+            if self.update_task:
+                task_status = self.update_task.get_status()
+                health_details.update(
+                    {
+                        "update_task_running": task_status["running"],
+                        "consecutive_failures": task_status["consecutive_failures"],
+                        "last_success_ago_seconds": task_status[
+                            "last_success_ago_seconds"
+                        ],
+                        "next_run_in_seconds": task_status["next_run_in_seconds"],
+                    }
+                )
+
+            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),
+            }
+
+    def get_pricing_service(self) -> PricingService | None:
+        """Get the pricing service instance."""
+        return self.service
+
+
+class PricingFactory(SystemPluginFactory):
+    """Factory for pricing plugin."""
+
+    def __init__(self) -> None:
+        """Initialize factory with manifest."""
+        # Create manifest with static declarations
+        manifest = PluginManifest(
+            name="pricing",
+            version="0.1.0",
+            description="Dynamic pricing plugin for AI model cost calculation",
+            is_provider=False,
+            config_class=PricingConfig,
+            provides=["pricing"],  # This plugin provides the pricing service
+        )
+
+        # Initialize with manifest
+        super().__init__(manifest)
+
+    def create_runtime(self) -> PricingRuntime:
+        """Create runtime instance."""
+        return PricingRuntime(self.manifest)
+
+
+# Export the factory instance
+factory = PricingFactory()

ccproxy/plugins/pricing/service.py

@@ -0,0 +1,191 @@
+"""Pricing service providing unified interface for pricing functionality."""
+
+from decimal import Decimal
+from typing import Any
+
+from ccproxy.core.logging import get_plugin_logger
+
+from .cache import PricingCache
+from .config import PricingConfig
+from .exceptions import (
+    ModelPricingNotFoundError,
+    PricingDataNotLoadedError,
+    PricingServiceDisabledError,
+)
+from .loader import PricingLoader
+from .models import ModelPricing, PricingData
+from .updater import PricingUpdater
+
+
+logger = get_plugin_logger(__name__)
+
+
+class PricingService:
+    """Main service interface for pricing functionality."""
+
+    def __init__(self, config: PricingConfig):
+        """Initialize pricing service with configuration."""
+        self.config = config
+        self.cache = PricingCache(config)
+        self.loader = PricingLoader()
+        self.updater = PricingUpdater(self.cache, config)
+        self._current_pricing: PricingData | None = None
+
+    async def initialize(self) -> None:
+        """Initialize the pricing service."""
+        if not self.config.enabled:
+            logger.info("pricing_service_disabled")
+            return
+
+        logger.debug("pricing_service_initializing")
+
+        # Force refresh on startup if configured
+        if self.config.force_refresh_on_startup:
+            await self.force_refresh_pricing()
+        else:
+            # Load current pricing data
+            await self.get_current_pricing()
+
+    async def get_current_pricing(
+        self, force_refresh: bool = False
+    ) -> PricingData | None:
+        """Get current pricing data."""
+        if not self.config.enabled:
+            return None
+
+        if force_refresh or self._current_pricing is None:
+            self._current_pricing = await self.updater.get_current_pricing(
+                force_refresh
+            )
+
+        return self._current_pricing
+
+    async def get_model_pricing(self, model_name: str) -> ModelPricing | None:
+        """Get pricing for specific model."""
+        pricing_data = await self.get_current_pricing()
+        if pricing_data is None:
+            return None
+
+        return pricing_data.get(model_name)
+
+    async def calculate_cost(
+        self,
+        model_name: str,
+        input_tokens: int = 0,
+        output_tokens: int = 0,
+        cache_read_tokens: int = 0,
+        cache_write_tokens: int = 0,
+    ) -> Decimal:
+        """Calculate cost for token usage.
+
+        Raises:
+            PricingServiceDisabledError: If pricing service is disabled
+            ModelPricingNotFoundError: If model pricing is not found
+        """
+        if not self.config.enabled:
+            raise PricingServiceDisabledError()
+
+        model_pricing = await self.get_model_pricing(model_name)
+        if model_pricing is None:
+            raise ModelPricingNotFoundError(model_name)
+
+        # Calculate cost per million tokens, then scale to actual tokens
+        total_cost = Decimal("0")
+
+        if input_tokens > 0:
+            total_cost += (model_pricing.input * input_tokens) / Decimal("1000000")
+
+        if output_tokens > 0:
+            total_cost += (model_pricing.output * output_tokens) / Decimal("1000000")
+
+        if cache_read_tokens > 0:
+            total_cost += (model_pricing.cache_read * cache_read_tokens) / Decimal(
+                "1000000"
+            )
+
+        if cache_write_tokens > 0:
+            total_cost += (model_pricing.cache_write * cache_write_tokens) / Decimal(
+                "1000000"
+            )
+
+        return total_cost
+
+    def calculate_cost_sync(
+        self,
+        model_name: str,
+        input_tokens: int = 0,
+        output_tokens: int = 0,
+        cache_read_tokens: int = 0,
+        cache_write_tokens: int = 0,
+    ) -> Decimal:
+        """Calculate cost synchronously using cached pricing data.
+
+        This method uses the cached pricing data and doesn't make any async calls,
+        making it safe to use in streaming contexts where we can't await.
+
+        Raises:
+            PricingServiceDisabledError: If pricing service is disabled
+            PricingDataNotLoadedError: If pricing data is not loaded yet
+            ModelPricingNotFoundError: If model pricing is not found
+        """
+        if not self.config.enabled:
+            raise PricingServiceDisabledError()
+
+        if self._current_pricing is None:
+            raise PricingDataNotLoadedError()
+
+        model_pricing = self._current_pricing.get(model_name)
+        if model_pricing is None:
+            raise ModelPricingNotFoundError(model_name)
+
+        # Calculate cost per million tokens, then scale to actual tokens
+        total_cost = Decimal("0")
+
+        if input_tokens > 0:
+            total_cost += (model_pricing.input * input_tokens) / Decimal("1000000")
+
+        if output_tokens > 0:
+            total_cost += (model_pricing.output * output_tokens) / Decimal("1000000")
+
+        if cache_read_tokens > 0:
+            total_cost += (model_pricing.cache_read * cache_read_tokens) / Decimal(
+                "1000000"
+            )
+
+        if cache_write_tokens > 0:
+            total_cost += (model_pricing.cache_write * cache_write_tokens) / Decimal(
+                "1000000"
+            )
+
+        return total_cost
+
+    async def force_refresh_pricing(self) -> bool:
+        """Force refresh of pricing data."""
+        if not self.config.enabled:
+            return False
+
+        success = await self.updater.force_refresh()
+        if success:
+            # Reload the current pricing data after successful refresh
+            self._current_pricing = await self.updater.get_current_pricing(
+                force_refresh=True
+            )
+            return True
+        return False
+
+    async def get_available_models(self) -> list[str]:
+        """Get list of available models with pricing."""
+        pricing_data = await self.get_current_pricing()
+        if pricing_data is None:
+            return []
+
+        return pricing_data.model_names()
+
+    def get_cache_info(self) -> dict[str, Any]:
+        """Get cache status information."""
+        return self.cache.get_cache_info()
+
+    async def clear_cache(self) -> bool:
+        """Clear pricing cache."""
+        self._current_pricing = None
+        return self.cache.clear_cache()

ccproxy/plugins/pricing/tasks.py

@@ -0,0 +1,300 @@
+"""Pricing plugin scheduled tasks."""
+
+import asyncio
+import contextlib
+import random
+import time
+from abc import ABC, abstractmethod
+from typing import Any
+
+from ccproxy.core.async_task_manager import create_managed_task
+from ccproxy.core.logging import get_plugin_logger
+
+from .service import PricingService
+
+
+logger = get_plugin_logger(__name__)
+
+
+class BaseScheduledTask(ABC):
+    """
+    Abstract base class for all scheduled tasks.
+
+    Provides common functionality for task lifecycle management, error handling,
+    and exponential backoff for failed executions.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        interval_seconds: float,
+        enabled: bool = True,
+        max_backoff_seconds: float = 300.0,
+        jitter_factor: float = 0.25,
+    ):
+        """
+        Initialize scheduled task.
+
+        Args:
+            name: Human-readable task name
+            interval_seconds: Interval between task executions in seconds
+            enabled: Whether the task is enabled
+            max_backoff_seconds: Maximum backoff delay for failed tasks
+            jitter_factor: Jitter factor for backoff randomization (0.0-1.0)
+        """
+        self.name = name
+        self.interval_seconds = max(1.0, interval_seconds)
+        self.enabled = enabled
+        self.max_backoff_seconds = max_backoff_seconds
+        self.jitter_factor = min(1.0, max(0.0, jitter_factor))
+
+        # Task state
+        self._task: asyncio.Task[None] | None = None
+        self._stop_event = asyncio.Event()
+        self._consecutive_failures = 0
+        self._last_success_time: float | None = None
+        self._next_run_time: float | None = None
+
+    @abstractmethod
+    async def run(self) -> bool:
+        """
+        Execute the task logic.
+
+        Returns:
+            True if task completed successfully, False otherwise
+        """
+
+    async def setup(self) -> None:  # noqa: B027
+        """
+        Optional setup hook called before the task starts running.
+
+        Override this method to perform any initialization required by the task.
+        """
+        pass
+
+    async def teardown(self) -> None:  # noqa: B027
+        """
+        Optional teardown hook called when the task stops.
+
+        Override this method to perform any cleanup required by the task.
+        """
+        pass
+
+    def _calculate_next_run_delay(self, failed: bool = False) -> float:
+        """Calculate delay until next task execution with exponential backoff."""
+        if not failed:
+            # Normal interval with jitter
+            base_delay = self.interval_seconds
+            jitter = random.uniform(-self.jitter_factor, self.jitter_factor)
+            return float(base_delay * (1 + jitter))
+
+        # Exponential backoff for failures
+        backoff_factor = min(2**self._consecutive_failures, 32)
+        backoff_delay = min(
+            self.interval_seconds * backoff_factor, self.max_backoff_seconds
+        )
+
+        # Add jitter to prevent thundering herd
+        jitter = random.uniform(-self.jitter_factor, self.jitter_factor)
+        return float(backoff_delay * (1 + jitter))
+
+    async def _run_with_error_handling(self) -> bool:
+        """Execute task with error handling and metrics."""
+        start_time = time.time()
+
+        try:
+            success = await self.run()
+
+            if success:
+                self._consecutive_failures = 0
+                self._last_success_time = start_time
+                logger.debug(
+                    "scheduled_task_success",
+                    task_name=self.name,
+                    duration=time.time() - start_time,
+                )
+            else:
+                self._consecutive_failures += 1
+                logger.warning(
+                    "scheduled_task_failed",
+                    task_name=self.name,
+                    consecutive_failures=self._consecutive_failures,
+                    duration=time.time() - start_time,
+                )
+
+            return success
+
+        except Exception as e:
+            self._consecutive_failures += 1
+            logger.error(
+                "scheduled_task_error",
+                task_name=self.name,
+                error=str(e),
+                error_type=type(e).__name__,
+                consecutive_failures=self._consecutive_failures,
+                duration=time.time() - start_time,
+                exc_info=e,
+            )
+            return False
+
+    async def _task_loop(self) -> None:
+        """Main task execution loop."""
+        logger.info("scheduled_task_starting", task_name=self.name)
+
+        try:
+            # Run setup
+            with contextlib.suppress(Exception):
+                await self.setup()
+
+            while not self._stop_event.is_set():
+                # Execute task
+                success = await self._run_with_error_handling()
+
+                # Calculate next run delay
+                delay = self._calculate_next_run_delay(failed=not success)
+                self._next_run_time = time.time() + delay
+
+                # Wait for next execution or stop event
+                try:
+                    await asyncio.wait_for(self._stop_event.wait(), timeout=delay)
+                    break  # Stop event was set
+                except TimeoutError:
+                    continue  # Time to run again
+
+        finally:
+            # Run teardown
+            with contextlib.suppress(Exception):
+                await self.teardown()
+
+            logger.info("scheduled_task_stopped", task_name=self.name)
+
+    async def start(self) -> None:
+        """Start the scheduled task."""
+        if not self.enabled:
+            logger.info("scheduled_task_disabled", task_name=self.name)
+            return
+
+        if self._task and not self._task.done():
+            logger.warning("scheduled_task_already_running", task_name=self.name)
+            return
+
+        self._stop_event.clear()
+        self._task = await create_managed_task(
+            self._task_loop(), name=f"scheduled_task_{self.name}"
+        )
+
+    async def stop(self, timeout: float = 10.0) -> None:
+        """Stop the scheduled task."""
+        if not self._task:
+            return
+
+        logger.info("scheduled_task_stopping", task_name=self.name)
+
+        # Signal stop
+        self._stop_event.set()
+
+        # Wait for task to complete
+        try:
+            await asyncio.wait_for(self._task, timeout=timeout)
+        except TimeoutError:
+            logger.warning(
+                "scheduled_task_stop_timeout", task_name=self.name, timeout=timeout
+            )
+            if not self._task.done():
+                self._task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await self._task
+
+        self._task = None
+
+    def is_running(self) -> bool:
+        """Check if task is currently running."""
+        return self._task is not None and not self._task.done()
+
+    def get_status(self) -> dict[str, Any]:
+        """Get current task status information."""
+        now = time.time()
+        return {
+            "name": self.name,
+            "enabled": self.enabled,
+            "running": self.is_running(),
+            "consecutive_failures": self._consecutive_failures,
+            "last_success_time": self._last_success_time,
+            "last_success_ago_seconds": (
+                now - self._last_success_time if self._last_success_time else None
+            ),
+            "next_run_time": self._next_run_time,
+            "next_run_in_seconds": (
+                self._next_run_time - now if self._next_run_time else None
+            ),
+            "interval_seconds": self.interval_seconds,
+        }
+
+
+class PricingCacheUpdateTask(BaseScheduledTask):
+    """Task for updating pricing cache periodically."""
+
+    def __init__(
+        self,
+        name: str,
+        interval_seconds: float,
+        pricing_service: PricingService,
+        enabled: bool = True,
+        force_refresh_on_startup: bool = False,
+    ):
+        """
+        Initialize pricing cache update task.
+
+        Args:
+            name: Task name
+            interval_seconds: Interval between pricing updates
+            pricing_service: Pricing service instance
+            enabled: Whether task is enabled
+            force_refresh_on_startup: Whether to force refresh on first run
+        """
+        super().__init__(
+            name=name,
+            interval_seconds=interval_seconds,
+            enabled=enabled,
+        )
+        self.pricing_service = pricing_service
+        self.force_refresh_on_startup = force_refresh_on_startup
+        self._first_run = True
+
+    async def run(self) -> bool:
+        """Execute pricing cache update."""
+        try:
+            if not self.pricing_service.config.enabled:
+                logger.debug("pricing_service_disabled", task_name=self.name)
+                return True  # Not a failure, just disabled
+
+            # Force refresh on first run if configured
+            force_refresh = self._first_run and self.force_refresh_on_startup
+            self._first_run = False
+
+            if force_refresh:
+                logger.info("pricing_update_force_refresh_startup", task_name=self.name)
+                success = await self.pricing_service.force_refresh_pricing()
+            else:
+                # Regular update check
+                pricing_data = await self.pricing_service.get_current_pricing(
+                    force_refresh=False
+                )
+                success = pricing_data is not None
+
+            if success:
+                logger.debug("pricing_update_success", task_name=self.name)
+            else:
+                logger.warning("pricing_update_failed", task_name=self.name)
+
+            return success
+
+        except Exception as e:
+            logger.error(
+                "pricing_update_task_error",
+                task_name=self.name,
+                error=str(e),
+                error_type=type(e).__name__,
+                exc_info=e,
+            )
+            return False