ccproxy-api 0.1.0__py3-none-any.whl

ccproxy/pricing/updater.py

@@ -0,0 +1,309 @@
+"""Pricing updater for managing periodic refresh of pricing data."""
+
+import asyncio
+from decimal import Decimal
+from typing import Any
+
+from structlog import get_logger
+
+from ccproxy.config.pricing import PricingSettings
+
+from .cache import PricingCache
+from .loader import PricingLoader
+from .models import PricingData
+
+
+logger = get_logger(__name__)
+
+
+class PricingUpdater:
+    """Manages periodic updates of pricing data."""
+
+    def __init__(
+        self,
+        cache: PricingCache,
+        settings: PricingSettings,
+    ) -> None:
+        """Initialize pricing updater.
+
+        Args:
+            cache: Pricing cache instance
+            settings: Pricing configuration settings
+        """
+        self.cache = cache
+        self.settings = settings
+        self._cached_pricing: PricingData | None = None
+        self._last_load_time: float = 0
+        self._last_file_check_time: float = 0
+        self._cached_file_mtime: float = 0
+
+    async def get_current_pricing(
+        self, force_refresh: bool = False
+    ) -> PricingData | None:
+        """Get current pricing data with automatic updates.
+
+        Args:
+            force_refresh: Force refresh even if cache is valid
+
+        Returns:
+            Current pricing data as PricingData model
+        """
+        import time
+
+        current_time = time.time()
+
+        # Return cached pricing if recent and not forced
+        if (
+            not force_refresh
+            and self._cached_pricing is not None
+            and (current_time - self._last_load_time) < self.settings.memory_cache_ttl
+        ):
+            # Only check file changes every 30 seconds to reduce I/O
+            if (current_time - self._last_file_check_time) > 30:
+                if self._has_cache_file_changed():
+                    logger.info("cache_file_changed")
+                    # File changed, need to reload
+                    pricing_data = await self._load_pricing_data()
+                    self._cached_pricing = pricing_data
+                    self._last_load_time = current_time
+                    return pricing_data
+                self._last_file_check_time = current_time
+
+            return self._cached_pricing
+
+        # Check if we need to refresh
+        should_refresh = force_refresh or (
+            self.settings.auto_update and not self.cache.is_cache_valid()
+        )
+
+        if should_refresh:
+            logger.info("pricing_refresh_start")
+            await self._refresh_pricing()
+
+        # Load pricing data
+        pricing_data = await self._load_pricing_data()
+
+        # Cache the result
+        self._cached_pricing = pricing_data
+        self._last_load_time = current_time
+        self._last_file_check_time = current_time
+
+        return pricing_data
+
+    def _has_cache_file_changed(self) -> bool:
+        """Check if the cache file has changed since last load.
+
+        Returns:
+            True if file has changed or doesn't exist
+        """
+        try:
+            if not self.cache.cache_file.exists():
+                return self._cached_file_mtime != 0  # File was deleted
+
+            current_mtime = self.cache.cache_file.stat().st_mtime
+            if current_mtime != self._cached_file_mtime:
+                self._cached_file_mtime = current_mtime
+                return True
+            return False
+        except OSError:
+            # If we can't check, assume it changed
+            return True
+
+    async def _refresh_pricing(self) -> bool:
+        """Refresh pricing data from external source.
+
+        Returns:
+            True if refresh was successful
+        """
+        try:
+            logger.info("pricing_refresh_start")
+
+            # Download fresh data
+            raw_data = await self.cache.download_pricing_data()
+            if raw_data is None:
+                logger.error("pricing_download_failed")
+                return False
+
+            # Save to cache
+            if not self.cache.save_to_cache(raw_data):
+                logger.error("cache_save_failed")
+                return False
+
+            logger.info("pricing_refresh_completed")
+            return True
+
+        except Exception as e:
+            logger.error("pricing_refresh_failed", error=str(e))
+            return False
+
+    async def _load_pricing_data(self) -> PricingData | None:
+        """Load pricing data from available sources.
+
+        Returns:
+            Pricing data as PricingData model
+        """
+        # Try to get data from cache or download
+        raw_data = await self.cache.get_pricing_data()
+
+        if raw_data is not None:
+            # Load and validate pricing data using Pydantic
+            pricing_data = PricingLoader.load_pricing_from_data(raw_data, verbose=False)
+
+            if pricing_data:
+                # Get cache info to display age
+                cache_info = self.cache.get_cache_info()
+                age_hours = cache_info.get("age_hours")
+
+                if age_hours is not None:
+                    logger.info(
+                        "pricing_loaded_from_external",
+                        model_count=len(pricing_data),
+                        cache_age_hours=round(age_hours, 2),
+                    )
+                else:
+                    logger.info(
+                        "pricing_loaded_from_external", model_count=len(pricing_data)
+                    )
+                return pricing_data
+            else:
+                logger.warning("external_pricing_validation_failed")
+
+        # Fallback to embedded pricing
+        if self.settings.fallback_to_embedded:
+            logger.info("using_embedded_pricing_fallback")
+            return self._get_embedded_pricing()
+        else:
+            logger.error("pricing_unavailable_no_fallback")
+            return None
+
+    def _get_embedded_pricing(self) -> PricingData:
+        """Get embedded (hardcoded) pricing as fallback.
+
+        Returns:
+            Embedded pricing data as PricingData model
+        """
+        # This is the current hardcoded pricing from CostCalculator
+        embedded_data = {
+            "claude-3-5-sonnet-20241022": {
+                "input": Decimal("3.00"),
+                "output": Decimal("15.00"),
+                "cache_read": Decimal("0.30"),
+                "cache_write": Decimal("3.75"),
+            },
+            "claude-3-5-haiku-20241022": {
+                "input": Decimal("0.25"),
+                "output": Decimal("1.25"),
+                "cache_read": Decimal("0.03"),
+                "cache_write": Decimal("0.30"),
+            },
+            "claude-3-opus-20240229": {
+                "input": Decimal("15.00"),
+                "output": Decimal("75.00"),
+                "cache_read": Decimal("1.50"),
+                "cache_write": Decimal("18.75"),
+            },
+            "claude-3-sonnet-20240229": {
+                "input": Decimal("3.00"),
+                "output": Decimal("15.00"),
+                "cache_read": Decimal("0.30"),
+                "cache_write": Decimal("3.75"),
+            },
+            "claude-3-haiku-20240307": {
+                "input": Decimal("0.25"),
+                "output": Decimal("1.25"),
+                "cache_read": Decimal("0.03"),
+                "cache_write": Decimal("0.30"),
+            },
+        }
+
+        # Create PricingData from embedded data
+        return PricingData.from_dict(embedded_data)
+
+    async def force_refresh(self) -> bool:
+        """Force a refresh of pricing data.
+
+        Returns:
+            True if refresh was successful
+        """
+        logger.info("pricing_force_refresh_start")
+
+        # Clear cached pricing
+        self._cached_pricing = None
+        self._last_load_time = 0
+
+        # Refresh from external source
+        success = await self._refresh_pricing()
+
+        if success:
+            # Reload pricing data
+            await self.get_current_pricing(force_refresh=True)
+
+        return success
+
+    def clear_cache(self) -> bool:
+        """Clear all cached pricing data.
+
+        Returns:
+            True if cache was cleared successfully
+        """
+        logger.info("pricing_cache_clear_start")
+
+        # Clear in-memory cache
+        self._cached_pricing = None
+        self._last_load_time = 0
+
+        # Clear file cache
+        return self.cache.clear_cache()
+
+    async def get_pricing_info(self) -> dict[str, Any]:
+        """Get information about current pricing state.
+
+        Returns:
+            Dictionary with pricing information
+        """
+        cache_info = self.cache.get_cache_info()
+
+        pricing_data = await self.get_current_pricing()
+
+        return {
+            "models_loaded": len(pricing_data) if pricing_data else 0,
+            "model_names": pricing_data.model_names() if pricing_data else [],
+            "auto_update": self.settings.auto_update,
+            "fallback_to_embedded": self.settings.fallback_to_embedded,
+            "has_cached_pricing": self._cached_pricing is not None,
+        }
+
+    async def validate_external_source(self) -> bool:
+        """Validate that external pricing source is accessible.
+
+        Returns:
+            True if external source is accessible and has valid data
+        """
+        try:
+            logger.debug("external_pricing_validation_start")
+
+            # Try to download data
+            raw_data = await self.cache.download_pricing_data(timeout=10)
+            if raw_data is None:
+                return False
+
+            # Try to parse Claude models
+            claude_models = PricingLoader.extract_claude_models(raw_data)
+            if not claude_models:
+                logger.warning("claude_models_not_found_in_external")
+                return False
+
+            # Try to load and validate using Pydantic
+            pricing_data = PricingLoader.load_pricing_from_data(raw_data, verbose=False)
+            if not pricing_data:
+                logger.warning("external_pricing_load_failed")
+                return False
+
+            logger.info(
+                "external_pricing_validation_completed", model_count=len(pricing_data)
+            )
+            return True
+
+        except Exception as e:
+            logger.error("external_pricing_validation_failed", error=str(e))
+            return False

ccproxy/scheduler/__init__.py

@@ -0,0 +1,39 @@
+"""
+Scheduler system for periodic tasks.
+
+This module provides a generic, extensible scheduler for managing periodic tasks
+in the CCProxy API. It provides a centralized system that supports:
+
+- Generic task scheduling with configurable intervals
+- Task registration and discovery via registry pattern
+- Graceful startup and shutdown with FastAPI integration
+- Error handling with exponential backoff
+- Structured logging and monitoring
+
+Key components:
+- Scheduler: Core scheduler engine for task management
+- BaseScheduledTask: Abstract base class for all scheduled tasks
+- TaskRegistry: Dynamic task registration and discovery system
+"""
+
+from .core import Scheduler
+from .registry import TaskRegistry, register_task
+from .tasks import (
+    BaseScheduledTask,
+    PricingCacheUpdateTask,
+    PushgatewayTask,
+    StatsPrintingTask,
+)
+
+
+# Task registration is now handled in manager.py during scheduler startup
+# to avoid side effects during module imports (e.g., CLI help display)
+
+__all__ = [
+    "Scheduler",
+    "TaskRegistry",
+    "BaseScheduledTask",
+    "PushgatewayTask",
+    "StatsPrintingTask",
+    "PricingCacheUpdateTask",
+]

ccproxy/scheduler/core.py

@@ -0,0 +1,335 @@
+"""Core scheduler for managing periodic tasks."""
+
+import asyncio
+from typing import Any
+
+import structlog
+
+from .exceptions import (
+    SchedulerError,
+    SchedulerShutdownError,
+    TaskNotFoundError,
+    TaskRegistrationError,
+)
+from .registry import TaskRegistry, get_task_registry
+from .tasks import BaseScheduledTask
+
+
+logger = structlog.get_logger(__name__)
+
+
+class Scheduler:
+    """
+    Scheduler for managing multiple periodic tasks.
+
+    Provides centralized management of scheduled tasks with:
+    - Dynamic task registration and configuration
+    - Graceful startup and shutdown
+    - Task monitoring and status reporting
+    - Error handling and recovery
+    """
+
+    def __init__(
+        self,
+        max_concurrent_tasks: int = 10,
+        graceful_shutdown_timeout: float = 30.0,
+        task_registry: TaskRegistry | None = None,
+    ):
+        """
+        Initialize the scheduler.
+
+        Args:
+            max_concurrent_tasks: Maximum number of tasks to run concurrently
+            graceful_shutdown_timeout: Timeout for graceful shutdown in seconds
+            task_registry: Task registry instance (uses global if None)
+        """
+        self.max_concurrent_tasks = max_concurrent_tasks
+        self.graceful_shutdown_timeout = graceful_shutdown_timeout
+        self.task_registry = task_registry or get_task_registry()
+
+        self._running = False
+        self._tasks: dict[str, BaseScheduledTask] = {}
+        self._semaphore: asyncio.Semaphore | None = None
+
+    async def start(self) -> None:
+        """Start the scheduler and all enabled tasks."""
+        if self._running:
+            logger.warning("scheduler_already_running")
+            return
+
+        self._running = True
+        self._semaphore = asyncio.Semaphore(self.max_concurrent_tasks)
+
+        logger.info(
+            "scheduler_starting",
+            max_concurrent_tasks=self.max_concurrent_tasks,
+            registered_tasks=self.task_registry.list_tasks(),
+        )
+
+        try:
+            # No automatic task creation - tasks must be explicitly added
+            logger.info(
+                "scheduler_started",
+                active_tasks=len(self._tasks),
+                running_tasks=[
+                    name for name, task in self._tasks.items() if task.is_running
+                ],
+            )
+        except Exception as e:
+            self._running = False
+            logger.error(
+                "scheduler_start_failed",
+                error=str(e),
+                error_type=type(e).__name__,
+            )
+            raise SchedulerError(f"Failed to start scheduler: {e}") from e
+
+    async def stop(self) -> None:
+        """Stop the scheduler and all running tasks."""
+        if not self._running:
+            return
+
+        self._running = False
+        logger.info("scheduler_stopping", active_tasks=len(self._tasks))
+
+        # Stop all tasks
+        stop_tasks = []
+        for task_name, task in self._tasks.items():
+            if task.is_running:
+                logger.debug("stopping_task", task_name=task_name)
+                stop_tasks.append(task.stop())
+
+        if stop_tasks:
+            try:
+                # Wait for all tasks to stop gracefully
+                await asyncio.wait_for(
+                    asyncio.gather(*stop_tasks, return_exceptions=True),
+                    timeout=self.graceful_shutdown_timeout,
+                )
+                logger.info("scheduler_stopped_gracefully")
+            except TimeoutError:
+                logger.warning(
+                    "scheduler_shutdown_timeout",
+                    timeout=self.graceful_shutdown_timeout,
+                )
+                # Tasks should have cancelled themselves, but log the issue
+                for task_name, task in self._tasks.items():
+                    if task.is_running:
+                        logger.warning(
+                            "task_still_running_after_shutdown", task_name=task_name
+                        )
+            except Exception as e:
+                logger.error(
+                    "scheduler_shutdown_error",
+                    error=str(e),
+                    error_type=type(e).__name__,
+                )
+                raise SchedulerShutdownError(
+                    f"Error during scheduler shutdown: {e}"
+                ) from e
+
+        self._tasks.clear()
+        logger.info("scheduler_stopped")
+
+    async def add_task(
+        self,
+        task_name: str,
+        task_type: str,
+        **task_kwargs: Any,
+    ) -> None:
+        """
+        Add and start a task.
+
+        Args:
+            task_name: Unique name for this task instance
+            task_type: Type of task (must be registered in task registry)
+            **task_kwargs: Additional arguments to pass to task constructor
+
+        Raises:
+            TaskRegistrationError: If task type is not registered
+            SchedulerError: If task name already exists or task creation fails
+        """
+        if task_name in self._tasks:
+            raise SchedulerError(f"Task '{task_name}' already exists")
+
+        if not self.task_registry.is_registered(task_type):
+            raise TaskRegistrationError(f"Task type '{task_type}' is not registered")
+
+        try:
+            # Get task class and create instance
+            task_class = self.task_registry.get(task_type)
+            task_instance = task_class(name=task_name, **task_kwargs)
+
+            # Add to our tasks dict
+            self._tasks[task_name] = task_instance
+
+            # Start the task if scheduler is running and task is enabled
+            if self._running and task_instance.enabled:
+                await task_instance.start()
+                logger.info(
+                    "task_added_and_started",
+                    task_name=task_name,
+                    task_type=task_type,
+                )
+            else:
+                logger.info(
+                    "task_added_not_started",
+                    task_name=task_name,
+                    task_type=task_type,
+                    scheduler_running=self._running,
+                    task_enabled=task_instance.enabled,
+                )
+
+        except Exception as e:
+            # Clean up if task was partially added
+            if task_name in self._tasks:
+                del self._tasks[task_name]
+
+            logger.error(
+                "task_add_failed",
+                task_name=task_name,
+                task_type=task_type,
+                error=str(e),
+                error_type=type(e).__name__,
+            )
+            raise SchedulerError(f"Failed to add task '{task_name}': {e}") from e
+
+    async def remove_task(self, task_name: str) -> None:
+        """
+        Remove and stop a task.
+
+        Args:
+            task_name: Name of task to remove
+
+        Raises:
+            TaskNotFoundError: If task does not exist
+        """
+        if task_name not in self._tasks:
+            raise TaskNotFoundError(f"Task '{task_name}' does not exist")
+
+        task = self._tasks[task_name]
+
+        try:
+            if task.is_running:
+                await task.stop()
+
+            del self._tasks[task_name]
+            logger.info("task_removed", task_name=task_name)
+
+        except Exception as e:
+            logger.error(
+                "task_remove_failed",
+                task_name=task_name,
+                error=str(e),
+                error_type=type(e).__name__,
+            )
+            raise SchedulerError(f"Failed to remove task '{task_name}': {e}") from e
+
+    def get_task(self, task_name: str) -> BaseScheduledTask:
+        """
+        Get a task instance by name.
+
+        Args:
+            task_name: Name of task to retrieve
+
+        Returns:
+            Task instance
+
+        Raises:
+            TaskNotFoundError: If task does not exist
+        """
+        if task_name not in self._tasks:
+            raise TaskNotFoundError(f"Task '{task_name}' does not exist")
+
+        return self._tasks[task_name]
+
+    def list_tasks(self) -> list[str]:
+        """
+        Get list of all task names.
+
+        Returns:
+            List of task names
+        """
+        return list(self._tasks.keys())
+
+    def get_task_status(self, task_name: str) -> dict[str, Any]:
+        """
+        Get status information for a specific task.
+
+        Args:
+            task_name: Name of task
+
+        Returns:
+            Task status dictionary
+
+        Raises:
+            TaskNotFoundError: If task does not exist
+        """
+        if task_name not in self._tasks:
+            raise TaskNotFoundError(f"Task '{task_name}' does not exist")
+
+        return self._tasks[task_name].get_status()
+
+    def get_scheduler_status(self) -> dict[str, Any]:
+        """
+        Get overall scheduler status information.
+
+        Returns:
+            Scheduler status dictionary
+        """
+        running_tasks = [name for name, task in self._tasks.items() if task.is_running]
+
+        return {
+            "running": self._running,
+            "total_tasks": len(self._tasks),
+            "running_tasks": len(running_tasks),
+            "max_concurrent_tasks": self.max_concurrent_tasks,
+            "graceful_shutdown_timeout": self.graceful_shutdown_timeout,
+            "task_names": list(self._tasks.keys()),
+            "running_task_names": running_tasks,
+            "registered_task_types": self.task_registry.list_tasks(),
+        }
+
+    @property
+    def is_running(self) -> bool:
+        """Check if the scheduler is running."""
+        return self._running
+
+    @property
+    def task_count(self) -> int:
+        """Get the number of managed tasks."""
+        return len(self._tasks)
+
+
+# Global scheduler instance
+_global_scheduler: Scheduler | None = None
+
+
+async def get_scheduler() -> Scheduler:
+    """
+    Get or create the global scheduler instance.
+
+    Returns:
+        Global Scheduler instance
+    """
+    global _global_scheduler
+
+    if _global_scheduler is None:
+        _global_scheduler = Scheduler()
+
+    return _global_scheduler
+
+
+async def start_scheduler() -> None:
+    """Start the global scheduler."""
+    scheduler = await get_scheduler()
+    await scheduler.start()
+
+
+async def stop_scheduler() -> None:
+    """Stop the global scheduler."""
+    global _global_scheduler
+
+    if _global_scheduler:
+        await _global_scheduler.stop()
+        _global_scheduler = None