llm-cost-guard 0.1.0__py3-none-any.whl

llm_cost_guard/tracker.py

@@ -0,0 +1,718 @@
+"""
+Main CostTracker class for LLM Cost Guard.
+"""
+
+import asyncio
+import functools
+import logging
+import time
+import threading
+from contextlib import contextmanager
+from datetime import datetime
+from typing import Any, Callable, Dict, List, Literal, Optional, TypeVar, Union
+
+from llm_cost_guard.backends import Backend, MemoryBackend, get_backend
+from llm_cost_guard.budget import Budget, BudgetAction, BudgetTracker
+from llm_cost_guard.exceptions import (
+    BudgetExceededError,
+    RateLimitExceededError,
+    TrackingUnavailableError,
+)
+from llm_cost_guard.models import CostRecord, CostReport, HealthStatus, ModelType, UsageData
+from llm_cost_guard.pricing.loader import PricingLoader, get_pricing_loader
+from llm_cost_guard.providers import detect_provider, get_provider
+from llm_cost_guard.rate_limit import RateLimit, RateLimiter
+from llm_cost_guard.span import Span, get_current_span
+
+logger = logging.getLogger(__name__)
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+class CostTracker:
+    """
+    Main entry point for cost tracking.
+
+    Provides decorator-based and context manager tracking for LLM API calls.
+    """
+
+    def __init__(
+        self,
+        budgets: Optional[List[Budget]] = None,
+        rate_limits: Optional[List[RateLimit]] = None,
+        backend: str = "memory",
+        auto_detect_provider: bool = True,
+        pricing_update: bool = True,
+        pricing_overrides: Optional[Dict[str, Dict[str, Any]]] = None,
+        on_tracking_failure: Literal["block", "allow", "fallback"] = "allow",
+        store_prompts: bool = False,
+        track_failed_calls: bool = True,
+        track_cache_savings: bool = True,
+        max_unique_tag_values: int = 1000,
+        budget_mode: Literal["local", "distributed"] = "local",
+        streaming_budget_mode: Literal["estimate", "actual"] = "actual",
+        streaming_max_output_estimate: int = 4096,
+        **backend_kwargs: Any,
+    ):
+        """
+        Initialize the CostTracker.
+
+        Args:
+            budgets: List of budget configurations
+            rate_limits: List of rate limit configurations
+            backend: Backend URL (memory, sqlite:///, postgresql://, redis://)
+            auto_detect_provider: Automatically detect provider from model name
+            pricing_update: Check for pricing updates on startup
+            pricing_overrides: Custom pricing for models (e.g., negotiated rates)
+            on_tracking_failure: Action when tracking fails (block/allow/fallback)
+            store_prompts: Store prompts in records (default: False for security)
+            track_failed_calls: Track costs for failed API calls
+            track_cache_savings: Track cache hit savings
+            max_unique_tag_values: Maximum unique values per tag key
+            budget_mode: Budget enforcement mode (local or distributed)
+            streaming_budget_mode: How to handle streaming budgets
+            streaming_max_output_estimate: Max output tokens to estimate for streaming
+        """
+        self._auto_detect_provider = auto_detect_provider
+        self._on_tracking_failure = on_tracking_failure
+        self._store_prompts = store_prompts
+        self._track_failed_calls = track_failed_calls
+        self._track_cache_savings = track_cache_savings
+        self._max_unique_tag_values = max_unique_tag_values
+        self._budget_mode = budget_mode
+        self._streaming_budget_mode = streaming_budget_mode
+        self._streaming_max_output_estimate = streaming_max_output_estimate
+
+        # Initialize backend
+        self._backend_url = backend
+        self._fallback_backend: Optional[MemoryBackend] = None
+        try:
+            self._backend: Backend = get_backend(backend, **backend_kwargs)
+        except Exception as e:
+            if on_tracking_failure == "block":
+                raise TrackingUnavailableError(f"Failed to initialize backend: {e}", backend)
+            elif on_tracking_failure == "fallback":
+                logger.warning(f"Failed to initialize backend {backend}, using memory fallback: {e}")
+                self._backend = MemoryBackend()
+                self._fallback_backend = self._backend
+            else:
+                logger.warning(f"Failed to initialize backend {backend}: {e}")
+                self._backend = MemoryBackend()
+
+        # Initialize pricing
+        self._pricing = PricingLoader(pricing_overrides=pricing_overrides)
+
+        # Initialize budget tracking
+        self._budget_tracker = BudgetTracker(budgets)
+
+        # Initialize rate limiting
+        self._rate_limiter = RateLimiter(rate_limits)
+
+        # Tag cardinality tracking
+        self._tag_values: Dict[str, set] = {}
+        self._tag_lock = threading.Lock()
+
+        # Last call tracking
+        self._last_record: Optional[CostRecord] = None
+        self._lock = threading.Lock()
+
+    def track(
+        self,
+        func: Optional[F] = None,
+        *,
+        tags: Optional[Dict[str, str]] = None,
+        streaming: bool = False,
+        provider: Optional[str] = None,
+        model: Optional[str] = None,
+    ) -> Union[F, Callable[[F], F]]:
+        """
+        Decorator to track LLM call costs.
+
+        Can be used with or without arguments:
+            @tracker.track
+            def my_call(): ...
+
+            @tracker.track(tags={"team": "search"})
+            def my_call(): ...
+
+        Args:
+            func: Function to decorate (when used without arguments)
+            tags: Tags for attribution
+            streaming: Whether the function returns a streaming response
+            provider: Override provider detection
+            model: Override model detection
+
+        Returns:
+            Decorated function
+        """
+
+        def decorator(f: F) -> F:
+            if asyncio.iscoroutinefunction(f):
+                return self._wrap_async(f, tags, streaming, provider, model)  # type: ignore
+            else:
+                return self._wrap_sync(f, tags, streaming, provider, model)  # type: ignore
+
+        if func is not None:
+            return decorator(func)
+        return decorator
+
+    def _wrap_sync(
+        self,
+        func: F,
+        tags: Optional[Dict[str, str]],
+        streaming: bool,
+        provider_override: Optional[str],
+        model_override: Optional[str],
+    ) -> F:
+        """Wrap a synchronous function for tracking."""
+
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            start_time = time.time()
+            success = True
+            error_type = None
+            response = None
+
+            try:
+                response = func(*args, **kwargs)
+                return response
+            except Exception as e:
+                success = False
+                error_type = type(e).__name__
+                raise
+            finally:
+                latency_ms = int((time.time() - start_time) * 1000)
+
+                if response is not None or (not success and self._track_failed_calls):
+                    try:
+                        self._record_call(
+                            response=response,
+                            tags=tags,
+                            success=success,
+                            error_type=error_type,
+                            latency_ms=latency_ms,
+                            provider_override=provider_override,
+                            model_override=model_override,
+                        )
+                    except Exception as e:
+                        self._handle_tracking_error(e)
+
+        return wrapper  # type: ignore
+
+    def _wrap_async(
+        self,
+        func: F,
+        tags: Optional[Dict[str, str]],
+        streaming: bool,
+        provider_override: Optional[str],
+        model_override: Optional[str],
+    ) -> F:
+        """Wrap an asynchronous function for tracking."""
+
+        @functools.wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> Any:
+            start_time = time.time()
+            success = True
+            error_type = None
+            response = None
+
+            try:
+                response = await func(*args, **kwargs)
+                return response
+            except Exception as e:
+                success = False
+                error_type = type(e).__name__
+                raise
+            finally:
+                latency_ms = int((time.time() - start_time) * 1000)
+
+                if response is not None or (not success and self._track_failed_calls):
+                    try:
+                        self._record_call(
+                            response=response,
+                            tags=tags,
+                            success=success,
+                            error_type=error_type,
+                            latency_ms=latency_ms,
+                            provider_override=provider_override,
+                            model_override=model_override,
+                        )
+                    except Exception as e:
+                        self._handle_tracking_error(e)
+
+        return wrapper  # type: ignore
+
+    @contextmanager
+    def track_context(
+        self,
+        tags: Optional[Dict[str, str]] = None,
+        provider: Optional[str] = None,
+        model: Optional[str] = None,
+    ):
+        """
+        Context manager for tracking LLM calls.
+
+        Usage:
+            with tracker.track_context(tags={"feature": "search"}):
+                response = openai.chat.completions.create(...)
+
+        Note: This context manager doesn't automatically extract usage from
+        responses. Use the decorator or manual recording for automatic tracking.
+        """
+        start_time = time.time()
+        tags = tags or {}
+
+        try:
+            yield
+        finally:
+            pass  # Context manager for grouping, actual tracking via decorator or record()
+
+    def span(
+        self,
+        name: str,
+        tags: Optional[Dict[str, str]] = None,
+    ) -> Span:
+        """
+        Create a tracking span for grouping multiple LLM calls.
+
+        Usage:
+            with tracker.span("rag_pipeline", tags={"user": "123"}) as span:
+                # Multiple LLM calls here
+                result = agent.run(query)
+                print(span.total_cost)
+
+        Args:
+            name: Name of the span
+            tags: Tags for the span
+
+        Returns:
+            Span context manager
+        """
+        return Span(name=name, tags=tags or {})
+
+    def record(
+        self,
+        provider: str,
+        model: str,
+        input_tokens: int,
+        output_tokens: int,
+        tags: Optional[Dict[str, str]] = None,
+        success: bool = True,
+        error_type: Optional[str] = None,
+        latency_ms: int = 0,
+        cached_tokens: int = 0,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> CostRecord:
+        """
+        Manually record an LLM call.
+
+        Use this for custom integrations or when automatic tracking isn't available.
+
+        Args:
+            provider: Provider name (openai, anthropic, bedrock)
+            model: Model name
+            input_tokens: Number of input tokens
+            output_tokens: Number of output tokens
+            tags: Attribution tags
+            success: Whether the call succeeded
+            error_type: Error type if call failed
+            latency_ms: Call latency in milliseconds
+            cached_tokens: Number of cached input tokens
+            metadata: Additional metadata (high-cardinality data)
+
+        Returns:
+            The created CostRecord
+        """
+        tags = tags or {}
+        metadata = metadata or {}
+
+        # Validate tag cardinality
+        self._check_tag_cardinality(tags)
+
+        # Calculate cost
+        input_cost, output_cost, total_cost = self._pricing.calculate_cost(
+            provider, model, input_tokens, output_tokens, cached_tokens
+        )
+
+        # Calculate cache savings
+        cache_savings = 0.0
+        if cached_tokens > 0 and self._track_cache_savings:
+            pricing = self._pricing.get_pricing(provider, model)
+            if pricing.cached_input_cost_per_1k is not None:
+                cache_savings = (cached_tokens / 1000) * (
+                    pricing.input_cost_per_1k - pricing.cached_input_cost_per_1k
+                )
+
+        # Check budgets
+        exceeded = self._budget_tracker.check_budget(total_cost, tags)
+        for budget, action in exceeded:
+            if action == BudgetAction.BLOCK:
+                raise BudgetExceededError(
+                    f"Budget '{budget.name}' would be exceeded",
+                    budget=budget,
+                    current=self._budget_tracker.get_spending(budget.name),
+                    limit=budget.limit,
+                )
+
+        # Check rate limits
+        rate_exceeded = self._rate_limiter.check(model=model, provider=provider, tags=tags)
+        if rate_exceeded:
+            limit, current, retry_after = rate_exceeded[0]
+            raise RateLimitExceededError(
+                f"Rate limit '{limit.name}' exceeded",
+                limit_name=limit.name,
+                current=current,
+                limit=limit.limit,
+                retry_after_seconds=retry_after,
+            )
+
+        # Create record
+        record = CostRecord(
+            timestamp=datetime.now(),
+            provider=provider,
+            model=model,
+            model_type=ModelType.CHAT,  # Default, can be enhanced
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            input_cost=input_cost,
+            output_cost=output_cost,
+            total_cost=total_cost,
+            latency_ms=latency_ms,
+            tags=tags,
+            metadata=metadata,
+            success=success,
+            error_type=error_type,
+            cached=cached_tokens > 0,
+            cache_savings=cache_savings,
+            span_id=get_current_span().span_id if get_current_span() else None,
+        )
+
+        # Save to backend
+        try:
+            self._backend.save_record(record)
+        except Exception as e:
+            self._handle_tracking_error(e)
+
+        # Record against budgets
+        self._budget_tracker.record_cost(total_cost, tags)
+
+        # Record rate limit usage
+        self._rate_limiter.record(model=model, provider=provider, tags=tags)
+
+        # Update last record
+        with self._lock:
+            self._last_record = record
+
+        # Record in current span if any
+        current_span = get_current_span()
+        if current_span:
+            current_span.record_call(
+                cost=total_cost,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
+                model=model,
+                record=record,
+            )
+
+        return record
+
+    def _record_call(
+        self,
+        response: Any,
+        tags: Optional[Dict[str, str]],
+        success: bool,
+        error_type: Optional[str],
+        latency_ms: int,
+        provider_override: Optional[str],
+        model_override: Optional[str],
+    ) -> Optional[CostRecord]:
+        """Record a call from a wrapped function."""
+        if response is None:
+            return None
+
+        # Detect provider and model
+        if model_override:
+            model = model_override
+            provider = provider_override or (
+                detect_provider(model) if self._auto_detect_provider else "unknown"
+            )
+        else:
+            # Try to extract from response
+            provider = provider_override or "openai"  # Default
+            model = "unknown"
+
+            # Try to detect and extract
+            try:
+                if self._auto_detect_provider:
+                    # Try OpenAI-style response
+                    if hasattr(response, "model"):
+                        model = response.model
+                        provider = detect_provider(model)
+                    elif isinstance(response, dict) and "model" in response:
+                        model = response["model"]
+                        provider = detect_provider(model)
+            except Exception:
+                pass
+
+        # Get provider handler
+        try:
+            provider_handler = get_provider(provider)
+        except ValueError:
+            logger.warning(f"Unknown provider {provider}, skipping cost tracking")
+            return None
+
+        # Extract usage
+        usage = provider_handler.extract_usage(response)
+        if model == "unknown":
+            model = provider_handler.extract_model(response)
+
+        # Record
+        return self.record(
+            provider=provider,
+            model=model,
+            input_tokens=usage.input_tokens,
+            output_tokens=usage.output_tokens,
+            tags=tags,
+            success=success,
+            error_type=error_type,
+            latency_ms=latency_ms,
+            cached_tokens=usage.cached_tokens,
+        )
+
+    def _handle_tracking_error(self, error: Exception) -> None:
+        """Handle errors during tracking based on configuration."""
+        if self._on_tracking_failure == "block":
+            raise TrackingUnavailableError(str(error), self._backend_url)
+        elif self._on_tracking_failure == "fallback":
+            logger.warning(f"Tracking error, using fallback: {error}")
+            if self._fallback_backend is None:
+                self._fallback_backend = MemoryBackend()
+        else:
+            logger.warning(f"Tracking error (allowing): {error}")
+
+    def _check_tag_cardinality(self, tags: Dict[str, str]) -> None:
+        """Check and track tag cardinality."""
+        with self._tag_lock:
+            for key, value in tags.items():
+                if key not in self._tag_values:
+                    self._tag_values[key] = set()
+
+                self._tag_values[key].add(value)
+
+                if len(self._tag_values[key]) > self._max_unique_tag_values:
+                    logger.warning(
+                        f"Tag '{key}' has exceeded cardinality limit "
+                        f"({len(self._tag_values[key])} > {self._max_unique_tag_values})"
+                    )
+
+    def last_call(self) -> Optional[CostRecord]:
+        """Get the last recorded call."""
+        with self._lock:
+            return self._last_record
+
+    def get_costs(
+        self,
+        start_date: Optional[str] = None,
+        end_date: Optional[str] = None,
+        tags: Optional[Dict[str, str]] = None,
+        group_by: Optional[List[str]] = None,
+    ) -> CostReport:
+        """
+        Query tracked costs.
+
+        Args:
+            start_date: Start date (ISO format)
+            end_date: End date (ISO format)
+            tags: Filter by tags
+            group_by: Group results by fields (provider, model, or tag keys)
+
+        Returns:
+            CostReport with aggregated data
+        """
+        start = datetime.fromisoformat(start_date) if start_date else None
+        end = datetime.fromisoformat(end_date) if end_date else None
+
+        return self._backend.get_report(
+            start_date=start,
+            end_date=end,
+            tags=tags,
+            group_by=group_by,
+        )
+
+    def daily_report(self) -> CostReport:
+        """Get a report for today."""
+        today = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0)
+        return self._backend.get_report(start_date=today)
+
+    def report_by_model(self, period: str = "day") -> CostReport:
+        """Get a report grouped by model."""
+        start = self._get_period_start(period)
+        return self._backend.get_report(start_date=start, group_by=["model"])
+
+    def trend_analysis(
+        self,
+        metric: str = "cost",
+        granularity: str = "hour",
+        last_n_days: int = 7,
+    ) -> Dict[str, Any]:
+        """Get trend analysis for a metric."""
+        # This is a simplified implementation
+        from datetime import timedelta
+
+        end = datetime.now()
+        start = end - timedelta(days=last_n_days)
+
+        records = self._backend.get_records(start_date=start, end_date=end)
+
+        # Group by time bucket
+        buckets: Dict[str, float] = {}
+        for record in records:
+            if granularity == "hour":
+                bucket_key = record.timestamp.strftime("%Y-%m-%d %H:00")
+            elif granularity == "day":
+                bucket_key = record.timestamp.strftime("%Y-%m-%d")
+            else:
+                bucket_key = record.timestamp.strftime("%Y-%m-%d %H:00")
+
+            if bucket_key not in buckets:
+                buckets[bucket_key] = 0.0
+
+            if metric == "cost":
+                buckets[bucket_key] += record.total_cost
+            elif metric == "tokens":
+                buckets[bucket_key] += record.input_tokens + record.output_tokens
+            elif metric == "calls":
+                buckets[bucket_key] += 1
+
+        return {
+            "metric": metric,
+            "granularity": granularity,
+            "start_date": start.isoformat(),
+            "end_date": end.isoformat(),
+            "data": buckets,
+        }
+
+    def _get_period_start(self, period: str) -> datetime:
+        """Get the start datetime for a period."""
+        from datetime import timedelta
+
+        now = datetime.now()
+
+        if period == "day":
+            return now.replace(hour=0, minute=0, second=0, microsecond=0)
+        elif period == "week":
+            start = now.replace(hour=0, minute=0, second=0, microsecond=0)
+            return start - timedelta(days=now.weekday())
+        elif period == "month":
+            return now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
+        else:
+            return now.replace(hour=0, minute=0, second=0, microsecond=0)
+
+    def to_dataframe(self):
+        """Export records to a pandas DataFrame."""
+        try:
+            import pandas as pd
+        except ImportError:
+            raise ImportError("pandas is required for DataFrame export. Install with: pip install pandas")
+
+        records = self._backend.get_records()
+
+        data = []
+        for r in records:
+            row = {
+                "timestamp": r.timestamp,
+                "provider": r.provider,
+                "model": r.model,
+                "input_tokens": r.input_tokens,
+                "output_tokens": r.output_tokens,
+                "total_cost": r.total_cost,
+                "latency_ms": r.latency_ms,
+                "success": r.success,
+                "cached": r.cached,
+            }
+            # Add tags as columns
+            for key, value in r.tags.items():
+                row[f"tag_{key}"] = value
+            data.append(row)
+
+        return pd.DataFrame(data)
+
+    def health_check(self) -> HealthStatus:
+        """Check tracker and backend health."""
+        errors = []
+
+        # Check backend
+        backend_connected = False
+        try:
+            backend_connected = self._backend.health_check()
+        except Exception as e:
+            errors.append(f"Backend health check failed: {e}")
+
+        # Check pricing freshness
+        pricing_fresh = not self._pricing.is_stale
+        if self._pricing.is_stale:
+            errors.append("Pricing data is stale")
+
+        # Get last record time
+        last_record_time = None
+        with self._lock:
+            if self._last_record:
+                last_record_time = self._last_record.timestamp
+
+        return HealthStatus(
+            healthy=backend_connected and pricing_fresh and len(errors) == 0,
+            backend_connected=backend_connected,
+            pricing_fresh=pricing_fresh,
+            last_record_time=last_record_time,
+            pending_records=0,
+            errors=errors,
+            pricing_version=str(self._pricing.pricing_version),
+            pricing_last_updated=self._pricing.last_updated,
+        )
+
+    def on_budget_warning(self, callback: Callable[[Budget, float], None]) -> None:
+        """Register a callback for budget warnings."""
+        self._budget_tracker.on_warning(callback)
+
+    def on_budget_exceeded(self, callback: Callable[[Budget], None]) -> None:
+        """Register a callback for budget exceeded events."""
+        self._budget_tracker.on_exceeded(callback)
+
+    def get_budget(self, name: str) -> Optional[Budget]:
+        """Get a budget by name."""
+        return self._budget_tracker.get_budget(name)
+
+    def get_budget_utilization(self, name: str) -> float:
+        """Get budget utilization percentage."""
+        return self._budget_tracker.get_utilization(name)
+
+    def reset_budget(self, name: Optional[str] = None) -> None:
+        """Reset a budget or all budgets."""
+        self._budget_tracker.reset(name)
+
+    @property
+    def pricing_last_updated(self) -> Optional[datetime]:
+        """Get when pricing was last updated."""
+        return self._pricing.last_updated
+
+    @property
+    def pricing_version(self) -> Dict[str, str]:
+        """Get pricing versions for all providers."""
+        return self._pricing.pricing_version
+
+    @property
+    def pricing_is_stale(self) -> bool:
+        """Check if pricing is stale."""
+        return self._pricing.is_stale
+
+    def close(self) -> None:
+        """Close the tracker and backend."""
+        self._backend.close()
+        if self._fallback_backend:
+            self._fallback_backend.close()
+
+    def __enter__(self) -> "CostTracker":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()