proxilion 0.0.1__py3-none-any.whl

proxilion/observability/cost_tracker.py

@@ -0,0 +1,866 @@
+"""
+Token usage and cost tracking for Proxilion.
+
+Tracks token usage and costs for every tool call and LLM interaction.
+Essential for budgeting, chargebacks, and understanding agent costs.
+
+Example:
+    >>> from proxilion.observability.cost_tracker import (
+    ...     CostTracker, BudgetPolicy, ModelPricing
+    ... )
+    >>>
+    >>> # Create tracker with budget limits
+    >>> tracker = CostTracker(
+    ...     budget_policy=BudgetPolicy(
+    ...         max_cost_per_request=1.00,
+    ...         max_cost_per_user_per_day=50.00,
+    ...     )
+    ... )
+    >>>
+    >>> # Record usage
+    >>> record = tracker.record_usage(
+    ...     model="claude-sonnet-4-20250514",
+    ...     input_tokens=1000,
+    ...     output_tokens=500,
+    ...     user_id="user_123",
+    ...     tool_name="database_query",
+    ... )
+    >>> print(f"Cost: ${record.cost_usd:.4f}")
+    >>>
+    >>> # Get summary
+    >>> summary = tracker.get_summary(user_id="user_123")
+    >>> print(f"Total cost: ${summary.total_cost:.4f}")
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import threading
+from collections import defaultdict
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timedelta, timezone
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ModelPricing:
+    """
+    Pricing information for an LLM model.
+
+    Prices are in USD per 1,000 tokens.
+
+    Attributes:
+        model_name: Display name for the model.
+        input_price_per_1k: Cost per 1,000 input tokens.
+        output_price_per_1k: Cost per 1,000 output tokens.
+        cache_read_price_per_1k: Cost per 1,000 cached input tokens (optional).
+        cache_write_price_per_1k: Cost per 1,000 tokens written to cache (optional).
+    """
+
+    model_name: str
+    input_price_per_1k: float
+    output_price_per_1k: float
+    cache_read_price_per_1k: float = 0.0
+    cache_write_price_per_1k: float = 0.0
+
+    def calculate_cost(
+        self,
+        input_tokens: int,
+        output_tokens: int,
+        cache_read_tokens: int = 0,
+        cache_write_tokens: int = 0,
+    ) -> float:
+        """
+        Calculate the cost for a given usage.
+
+        Args:
+            input_tokens: Number of input tokens.
+            output_tokens: Number of output tokens.
+            cache_read_tokens: Number of cached input tokens read.
+            cache_write_tokens: Number of tokens written to cache.
+
+        Returns:
+            Total cost in USD.
+        """
+        input_cost = (input_tokens / 1000) * self.input_price_per_1k
+        output_cost = (output_tokens / 1000) * self.output_price_per_1k
+        cache_read_cost = (cache_read_tokens / 1000) * self.cache_read_price_per_1k
+        cache_write_cost = (cache_write_tokens / 1000) * self.cache_write_price_per_1k
+
+        return input_cost + output_cost + cache_read_cost + cache_write_cost
+
+
+# Default pricing for popular models (as of early 2026)
+DEFAULT_PRICING: dict[str, ModelPricing] = {
+    # Anthropic Claude models
+    "claude-opus-4-5-20251101": ModelPricing(
+        model_name="Claude Opus 4.5",
+        input_price_per_1k=0.015,
+        output_price_per_1k=0.075,
+        cache_read_price_per_1k=0.00375,
+        cache_write_price_per_1k=0.01875,
+    ),
+    "claude-sonnet-4-20250514": ModelPricing(
+        model_name="Claude Sonnet 4",
+        input_price_per_1k=0.003,
+        output_price_per_1k=0.015,
+        cache_read_price_per_1k=0.0006,
+        cache_write_price_per_1k=0.00375,
+    ),
+    "claude-3-5-sonnet-20241022": ModelPricing(
+        model_name="Claude 3.5 Sonnet",
+        input_price_per_1k=0.003,
+        output_price_per_1k=0.015,
+        cache_read_price_per_1k=0.0003,
+        cache_write_price_per_1k=0.00375,
+    ),
+    "claude-3-5-haiku-20241022": ModelPricing(
+        model_name="Claude 3.5 Haiku",
+        input_price_per_1k=0.001,
+        output_price_per_1k=0.005,
+        cache_read_price_per_1k=0.0001,
+        cache_write_price_per_1k=0.00125,
+    ),
+    # OpenAI models
+    "gpt-4o": ModelPricing(
+        model_name="GPT-4o",
+        input_price_per_1k=0.0025,
+        output_price_per_1k=0.01,
+        cache_read_price_per_1k=0.00125,
+    ),
+    "gpt-4o-mini": ModelPricing(
+        model_name="GPT-4o Mini",
+        input_price_per_1k=0.00015,
+        output_price_per_1k=0.0006,
+        cache_read_price_per_1k=0.000075,
+    ),
+    "gpt-4-turbo": ModelPricing(
+        model_name="GPT-4 Turbo",
+        input_price_per_1k=0.01,
+        output_price_per_1k=0.03,
+    ),
+    "gpt-3.5-turbo": ModelPricing(
+        model_name="GPT-3.5 Turbo",
+        input_price_per_1k=0.0005,
+        output_price_per_1k=0.0015,
+    ),
+    # Google models
+    "gemini-1.5-pro": ModelPricing(
+        model_name="Gemini 1.5 Pro",
+        input_price_per_1k=0.00125,
+        output_price_per_1k=0.005,
+        cache_read_price_per_1k=0.000315,
+    ),
+    "gemini-1.5-flash": ModelPricing(
+        model_name="Gemini 1.5 Flash",
+        input_price_per_1k=0.000075,
+        output_price_per_1k=0.0003,
+        cache_read_price_per_1k=0.00001875,
+    ),
+    "gemini-2.0-flash": ModelPricing(
+        model_name="Gemini 2.0 Flash",
+        input_price_per_1k=0.0001,
+        output_price_per_1k=0.0004,
+    ),
+}
+
+
+@dataclass
+class UsageRecord:
+    """
+    Record of a single usage event.
+
+    Attributes:
+        timestamp: When the usage occurred.
+        model: Model identifier.
+        input_tokens: Number of input tokens.
+        output_tokens: Number of output tokens.
+        cache_read_tokens: Number of cached tokens read.
+        cache_write_tokens: Number of tokens written to cache.
+        tool_name: Tool that triggered the usage (if any).
+        user_id: User who incurred the usage.
+        cost_usd: Calculated cost in USD.
+        request_id: Optional request identifier.
+        metadata: Additional metadata.
+    """
+
+    timestamp: datetime
+    model: str
+    input_tokens: int
+    output_tokens: int
+    cache_read_tokens: int = 0
+    cache_write_tokens: int = 0
+    tool_name: str | None = None
+    user_id: str | None = None
+    cost_usd: float = 0.0
+    request_id: str | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def total_tokens(self) -> int:
+        """Total tokens (input + output)."""
+        return self.input_tokens + self.output_tokens
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary."""
+        result = asdict(self)
+        result["timestamp"] = self.timestamp.isoformat()
+        return result
+
+    def to_json(self) -> str:
+        """Convert to JSON string."""
+        return json.dumps(self.to_dict())
+
+
+@dataclass
+class CostSummary:
+    """
+    Summary of costs over a period.
+
+    Attributes:
+        total_cost: Total cost in USD.
+        total_input_tokens: Total input tokens.
+        total_output_tokens: Total output tokens.
+        total_cache_tokens: Total cache read tokens.
+        record_count: Number of usage records.
+        by_model: Cost breakdown by model.
+        by_user: Cost breakdown by user.
+        by_tool: Cost breakdown by tool.
+        start_time: Start of the summary period.
+        end_time: End of the summary period.
+    """
+
+    total_cost: float = 0.0
+    total_input_tokens: int = 0
+    total_output_tokens: int = 0
+    total_cache_tokens: int = 0
+    record_count: int = 0
+    by_model: dict[str, float] = field(default_factory=dict)
+    by_user: dict[str, float] = field(default_factory=dict)
+    by_tool: dict[str, float] = field(default_factory=dict)
+    start_time: datetime | None = None
+    end_time: datetime | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "total_cost": self.total_cost,
+            "total_input_tokens": self.total_input_tokens,
+            "total_output_tokens": self.total_output_tokens,
+            "total_cache_tokens": self.total_cache_tokens,
+            "record_count": self.record_count,
+            "by_model": self.by_model,
+            "by_user": self.by_user,
+            "by_tool": self.by_tool,
+            "start_time": self.start_time.isoformat() if self.start_time else None,
+            "end_time": self.end_time.isoformat() if self.end_time else None,
+        }
+
+
+@dataclass
+class BudgetPolicy:
+    """
+    Budget limits for cost control.
+
+    Attributes:
+        max_cost_per_request: Maximum cost for a single request.
+        max_cost_per_user_per_day: Maximum daily cost per user.
+        max_cost_per_user_per_hour: Maximum hourly cost per user.
+        max_tokens_per_request: Maximum tokens per request.
+        max_org_cost_per_day: Maximum daily cost for the organization.
+        max_org_cost_per_month: Maximum monthly cost for the organization.
+        warn_at_percentage: Percentage at which to warn (0.0 to 1.0).
+    """
+
+    max_cost_per_request: float | None = None
+    max_cost_per_user_per_day: float | None = None
+    max_cost_per_user_per_hour: float | None = None
+    max_tokens_per_request: int | None = None
+    max_org_cost_per_day: float | None = None
+    max_org_cost_per_month: float | None = None
+    warn_at_percentage: float = 0.8
+
+
+class CostTracker:
+    """
+    Tracks token usage and costs across all operations.
+
+    The CostTracker maintains a history of usage records and can
+    calculate costs, check budgets, and provide summaries.
+
+    Example:
+        >>> tracker = CostTracker()
+        >>>
+        >>> # Record usage
+        >>> record = tracker.record_usage(
+        ...     model="claude-sonnet-4-20250514",
+        ...     input_tokens=1000,
+        ...     output_tokens=500,
+        ...     user_id="user_123",
+        ... )
+        >>>
+        >>> # Check budget before expensive operation
+        >>> allowed, reason = tracker.check_budget("user_123", estimated_tokens=10000)
+        >>> if not allowed:
+        ...     print(f"Budget issue: {reason}")
+    """
+
+    def __init__(
+        self,
+        pricing: dict[str, ModelPricing] | None = None,
+        budget_policy: BudgetPolicy | None = None,
+        max_records: int = 100000,
+        default_model: str = "claude-sonnet-4-20250514",
+    ) -> None:
+        """
+        Initialize the cost tracker.
+
+        Args:
+            pricing: Model pricing information. Defaults to DEFAULT_PRICING.
+            budget_policy: Budget limits to enforce.
+            max_records: Maximum records to keep in memory.
+            default_model: Default model for cost estimation.
+        """
+        self._lock = threading.RLock()
+        self._pricing = dict(DEFAULT_PRICING)
+        if pricing:
+            self._pricing.update(pricing)
+
+        self._budget_policy = budget_policy
+        self._max_records = max_records
+        self._default_model = default_model
+
+        # Storage
+        self._records: list[UsageRecord] = []
+        self._user_daily_spend: dict[str, dict[str, float]] = defaultdict(
+            lambda: defaultdict(float)
+        )
+        self._user_hourly_spend: dict[str, dict[str, float]] = defaultdict(
+            lambda: defaultdict(float)
+        )
+        self._org_daily_spend: dict[str, float] = defaultdict(float)
+        self._org_monthly_spend: dict[str, float] = defaultdict(float)
+
+    def set_pricing(self, model: str, pricing: ModelPricing) -> None:
+        """
+        Set or update pricing for a model.
+
+        Args:
+            model: Model identifier.
+            pricing: Pricing information.
+        """
+        with self._lock:
+            self._pricing[model] = pricing
+
+    def get_pricing(self, model: str) -> ModelPricing | None:
+        """
+        Get pricing for a model.
+
+        Args:
+            model: Model identifier.
+
+        Returns:
+            ModelPricing if found, None otherwise.
+        """
+        return self._pricing.get(model)
+
+    def record_usage(
+        self,
+        model: str,
+        input_tokens: int,
+        output_tokens: int,
+        cache_read_tokens: int = 0,
+        cache_write_tokens: int = 0,
+        tool_name: str | None = None,
+        user_id: str | None = None,
+        request_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        timestamp: datetime | None = None,
+    ) -> UsageRecord:
+        """
+        Record a usage event and calculate cost.
+
+        Args:
+            model: Model identifier.
+            input_tokens: Number of input tokens.
+            output_tokens: Number of output tokens.
+            cache_read_tokens: Number of cached tokens read.
+            cache_write_tokens: Number of tokens written to cache.
+            tool_name: Tool that triggered the usage.
+            user_id: User who incurred the usage.
+            request_id: Request identifier.
+            metadata: Additional metadata.
+            timestamp: Event timestamp (defaults to now).
+
+        Returns:
+            The created UsageRecord.
+        """
+        if timestamp is None:
+            timestamp = datetime.now(timezone.utc)
+
+        # Calculate cost
+        pricing = self._pricing.get(model)
+        if pricing:
+            cost = pricing.calculate_cost(
+                input_tokens, output_tokens, cache_read_tokens, cache_write_tokens
+            )
+        else:
+            # Use default model pricing as fallback
+            default_pricing = self._pricing.get(self._default_model)
+            if default_pricing:
+                cost = default_pricing.calculate_cost(
+                    input_tokens, output_tokens, cache_read_tokens, cache_write_tokens
+                )
+                logger.warning(
+                    f"Unknown model '{model}', using default pricing from '{self._default_model}'"
+                )
+            else:
+                cost = 0.0
+                logger.warning(f"Unknown model '{model}' and no default pricing available")
+
+        record = UsageRecord(
+            timestamp=timestamp,
+            model=model,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cache_read_tokens=cache_read_tokens,
+            cache_write_tokens=cache_write_tokens,
+            tool_name=tool_name,
+            user_id=user_id,
+            cost_usd=cost,
+            request_id=request_id,
+            metadata=metadata or {},
+        )
+
+        with self._lock:
+            self._records.append(record)
+
+            # Trim old records if needed
+            if len(self._records) > self._max_records:
+                self._records = self._records[-self._max_records:]
+
+            # Update spend tracking
+            if user_id:
+                day_key = timestamp.strftime("%Y-%m-%d")
+                hour_key = timestamp.strftime("%Y-%m-%d-%H")
+                self._user_daily_spend[user_id][day_key] += cost
+                self._user_hourly_spend[user_id][hour_key] += cost
+
+            # Track org-wide spend
+            day_key = timestamp.strftime("%Y-%m-%d")
+            month_key = timestamp.strftime("%Y-%m")
+            self._org_daily_spend[day_key] += cost
+            self._org_monthly_spend[month_key] += cost
+
+        logger.debug(
+            f"Recorded usage: model={model}, tokens={input_tokens}+{output_tokens}, "
+            f"cost=${cost:.4f}, user={user_id}"
+        )
+
+        return record
+
+    def estimate_cost(
+        self,
+        model: str | None = None,
+        input_tokens: int = 0,
+        output_tokens: int = 0,
+        cache_read_tokens: int = 0,
+    ) -> float:
+        """
+        Estimate cost for a potential request.
+
+        Args:
+            model: Model to use (defaults to default_model).
+            input_tokens: Expected input tokens.
+            output_tokens: Expected output tokens.
+            cache_read_tokens: Expected cached tokens.
+
+        Returns:
+            Estimated cost in USD.
+        """
+        model = model or self._default_model
+        pricing = self._pricing.get(model)
+
+        if pricing:
+            return pricing.calculate_cost(input_tokens, output_tokens, cache_read_tokens)
+
+        return 0.0
+
+    def check_budget(
+        self,
+        user_id: str,
+        estimated_cost: float = 0.0,
+        estimated_tokens: int = 0,
+    ) -> tuple[bool, str | None]:
+        """
+        Check if a request would exceed budget limits.
+
+        Args:
+            user_id: User making the request.
+            estimated_cost: Estimated cost of the request.
+            estimated_tokens: Estimated tokens for the request.
+
+        Returns:
+            Tuple of (allowed, reason). If not allowed, reason explains why.
+        """
+        if self._budget_policy is None:
+            return True, None
+
+        policy = self._budget_policy
+        now = datetime.now(timezone.utc)
+
+        # If estimated_cost not provided, estimate from tokens
+        if estimated_cost == 0.0 and estimated_tokens > 0:
+            # Assume 50/50 split between input/output for estimation
+            estimated_cost = self.estimate_cost(
+                input_tokens=estimated_tokens // 2,
+                output_tokens=estimated_tokens // 2,
+            )
+
+        # Check per-request cost limit
+        if policy.max_cost_per_request is not None:
+            if estimated_cost > policy.max_cost_per_request:
+                return False, (
+                    f"Request would exceed per-request budget: "
+                    f"${estimated_cost:.4f} > ${policy.max_cost_per_request:.4f}"
+                )
+
+        # Check per-request token limit
+        if policy.max_tokens_per_request is not None:
+            if estimated_tokens > policy.max_tokens_per_request:
+                return False, (
+                    f"Request would exceed token limit: "
+                    f"{estimated_tokens} > {policy.max_tokens_per_request}"
+                )
+
+        # Check user daily limit
+        if policy.max_cost_per_user_per_day is not None:
+            daily_spend = self.get_user_spend(user_id, timedelta(days=1))
+            if daily_spend + estimated_cost > policy.max_cost_per_user_per_day:
+                return False, (
+                    f"User daily budget exceeded: "
+                    f"${daily_spend + estimated_cost:.4f} > ${policy.max_cost_per_user_per_day:.4f}"
+                )
+
+        # Check user hourly limit
+        if policy.max_cost_per_user_per_hour is not None:
+            hourly_spend = self.get_user_spend(user_id, timedelta(hours=1))
+            if hourly_spend + estimated_cost > policy.max_cost_per_user_per_hour:
+                return False, (
+                    f"User hourly budget exceeded: ${hourly_spend + estimated_cost:.4f}"
+                    f" > ${policy.max_cost_per_user_per_hour:.4f}"
+                )
+
+        # Check org daily limit
+        if policy.max_org_cost_per_day is not None:
+            day_key = now.strftime("%Y-%m-%d")
+            org_daily = self._org_daily_spend.get(day_key, 0.0)
+            if org_daily + estimated_cost > policy.max_org_cost_per_day:
+                return False, (
+                    f"Organization daily budget exceeded: "
+                    f"${org_daily + estimated_cost:.4f} > ${policy.max_org_cost_per_day:.4f}"
+                )
+
+        # Check org monthly limit
+        if policy.max_org_cost_per_month is not None:
+            month_key = now.strftime("%Y-%m")
+            org_monthly = self._org_monthly_spend.get(month_key, 0.0)
+            if org_monthly + estimated_cost > policy.max_org_cost_per_month:
+                return False, (
+                    f"Organization monthly budget exceeded: "
+                    f"${org_monthly + estimated_cost:.4f} > ${policy.max_org_cost_per_month:.4f}"
+                )
+
+        return True, None
+
+    def get_user_spend(self, user_id: str, period: timedelta) -> float:
+        """
+        Get total spend for a user over a period.
+
+        Args:
+            user_id: User to check.
+            period: Time period to check.
+
+        Returns:
+            Total spend in USD.
+        """
+        with self._lock:
+            now = datetime.now(timezone.utc)
+            cutoff = now - period
+
+            total = 0.0
+            for record in reversed(self._records):
+                if record.timestamp < cutoff:
+                    break
+                if record.user_id == user_id:
+                    total += record.cost_usd
+
+            return total
+
+    def get_org_spend(self, period: timedelta) -> float:
+        """
+        Get total organization spend over a period.
+
+        Args:
+            period: Time period to check.
+
+        Returns:
+            Total spend in USD.
+        """
+        with self._lock:
+            now = datetime.now(timezone.utc)
+            cutoff = now - period
+
+            total = 0.0
+            for record in reversed(self._records):
+                if record.timestamp < cutoff:
+                    break
+                total += record.cost_usd
+
+            return total
+
+    def get_summary(
+        self,
+        start: datetime | None = None,
+        end: datetime | None = None,
+        user_id: str | None = None,
+        model: str | None = None,
+        tool_name: str | None = None,
+    ) -> CostSummary:
+        """
+        Get a cost summary for the specified period and filters.
+
+        Args:
+            start: Start of period (defaults to all time).
+            end: End of period (defaults to now).
+            user_id: Filter by user.
+            model: Filter by model.
+            tool_name: Filter by tool.
+
+        Returns:
+            CostSummary with aggregated data.
+        """
+        with self._lock:
+            summary = CostSummary(
+                start_time=start,
+                end_time=end or datetime.now(timezone.utc),
+            )
+
+            by_model: dict[str, float] = defaultdict(float)
+            by_user: dict[str, float] = defaultdict(float)
+            by_tool: dict[str, float] = defaultdict(float)
+
+            for record in self._records:
+                # Apply time filters
+                if start and record.timestamp < start:
+                    continue
+                if end and record.timestamp > end:
+                    continue
+
+                # Apply entity filters
+                if user_id and record.user_id != user_id:
+                    continue
+                if model and record.model != model:
+                    continue
+                if tool_name and record.tool_name != tool_name:
+                    continue
+
+                # Aggregate
+                summary.total_cost += record.cost_usd
+                summary.total_input_tokens += record.input_tokens
+                summary.total_output_tokens += record.output_tokens
+                summary.total_cache_tokens += record.cache_read_tokens
+                summary.record_count += 1
+
+                by_model[record.model] += record.cost_usd
+                if record.user_id:
+                    by_user[record.user_id] += record.cost_usd
+                if record.tool_name:
+                    by_tool[record.tool_name] += record.cost_usd
+
+            summary.by_model = dict(by_model)
+            summary.by_user = dict(by_user)
+            summary.by_tool = dict(by_tool)
+
+            return summary
+
+    def get_records(
+        self,
+        start: datetime | None = None,
+        end: datetime | None = None,
+        user_id: str | None = None,
+        limit: int | None = None,
+    ) -> list[UsageRecord]:
+        """
+        Get usage records with optional filters.
+
+        Args:
+            start: Start of period.
+            end: End of period.
+            user_id: Filter by user.
+            limit: Maximum records to return.
+
+        Returns:
+            List of UsageRecords, most recent first.
+        """
+        with self._lock:
+            result = []
+
+            for record in reversed(self._records):
+                if start and record.timestamp < start:
+                    continue
+                if end and record.timestamp > end:
+                    continue
+                if user_id and record.user_id != user_id:
+                    continue
+
+                result.append(record)
+
+                if limit and len(result) >= limit:
+                    break
+
+            return result
+
+    def export_records(
+        self,
+        format: str = "json",
+        start: datetime | None = None,
+        end: datetime | None = None,
+    ) -> str:
+        """
+        Export records to a string format.
+
+        Args:
+            format: Output format ("json" or "jsonl").
+            start: Start of period.
+            end: End of period.
+
+        Returns:
+            Exported records as string.
+        """
+        records = self.get_records(start=start, end=end)
+
+        if format == "jsonl":
+            return "\n".join(r.to_json() for r in records)
+        else:
+            return json.dumps([r.to_dict() for r in records], indent=2)
+
+    def clear_records(
+        self,
+        before: datetime | None = None,
+        user_id: str | None = None,
+    ) -> int:
+        """
+        Clear usage records.
+
+        Args:
+            before: Clear records before this time.
+            user_id: Clear only records for this user.
+
+        Returns:
+            Number of records cleared.
+        """
+        with self._lock:
+            original_count = len(self._records)
+
+            if before is None and user_id is None:
+                self._records.clear()
+                self._user_daily_spend.clear()
+                self._user_hourly_spend.clear()
+                self._org_daily_spend.clear()
+                self._org_monthly_spend.clear()
+                return original_count
+
+            self._records = [
+                r for r in self._records
+                if not (
+                    (before is None or r.timestamp >= before) and
+                    (user_id is None or r.user_id == user_id)
+                )
+            ]
+
+            return original_count - len(self._records)
+
+    def set_budget_policy(self, policy: BudgetPolicy | None) -> None:
+        """
+        Set or update the budget policy.
+
+        Args:
+            policy: New budget policy, or None to disable.
+        """
+        self._budget_policy = policy
+
+    def get_budget_policy(self) -> BudgetPolicy | None:
+        """Get the current budget policy."""
+        return self._budget_policy
+
+    def get_budget_status(self, user_id: str) -> dict[str, Any]:
+        """
+        Get current budget status for a user.
+
+        Args:
+            user_id: User to check.
+
+        Returns:
+            Dictionary with budget status information.
+        """
+        if self._budget_policy is None:
+            return {"policy_active": False}
+
+        policy = self._budget_policy
+        now = datetime.now(timezone.utc)
+
+        status: dict[str, Any] = {"policy_active": True}
+
+        if policy.max_cost_per_user_per_day is not None:
+            daily_spend = self.get_user_spend(user_id, timedelta(days=1))
+            status["daily"] = {
+                "spent": daily_spend,
+                "limit": policy.max_cost_per_user_per_day,
+                "remaining": max(0, policy.max_cost_per_user_per_day - daily_spend),
+                "percentage": daily_spend / policy.max_cost_per_user_per_day,
+            }
+
+        if policy.max_cost_per_user_per_hour is not None:
+            hourly_spend = self.get_user_spend(user_id, timedelta(hours=1))
+            status["hourly"] = {
+                "spent": hourly_spend,
+                "limit": policy.max_cost_per_user_per_hour,
+                "remaining": max(0, policy.max_cost_per_user_per_hour - hourly_spend),
+                "percentage": hourly_spend / policy.max_cost_per_user_per_hour,
+            }
+
+        if policy.max_org_cost_per_day is not None:
+            day_key = now.strftime("%Y-%m-%d")
+            org_daily = self._org_daily_spend.get(day_key, 0.0)
+            status["org_daily"] = {
+                "spent": org_daily,
+                "limit": policy.max_org_cost_per_day,
+                "remaining": max(0, policy.max_org_cost_per_day - org_daily),
+                "percentage": org_daily / policy.max_org_cost_per_day,
+            }
+
+        return status
+
+
+def create_cost_tracker(
+    budget_policy: BudgetPolicy | None = None,
+    custom_pricing: dict[str, ModelPricing] | None = None,
+) -> CostTracker:
+    """
+    Factory function to create a CostTracker.
+
+    Args:
+        budget_policy: Optional budget limits.
+        custom_pricing: Additional model pricing.
+
+    Returns:
+        Configured CostTracker instance.
+    """
+    return CostTracker(
+        pricing=custom_pricing,
+        budget_policy=budget_policy,
+    )