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

llm_cost_guard/__init__.py

@@ -0,0 +1,39 @@
+"""
+LLM Cost Guard - Real-time cost tracking, budget enforcement, and usage analytics for LLM applications.
+"""
+
+from llm_cost_guard.tracker import CostTracker
+from llm_cost_guard.budget import Budget, BudgetAction
+from llm_cost_guard.rate_limit import RateLimit
+from llm_cost_guard.span import Span
+from llm_cost_guard.models import CostRecord, CostReport, HealthStatus
+from llm_cost_guard.exceptions import (
+    LLMCostGuardError,
+    BudgetExceededError,
+    PricingNotFoundError,
+    TokenCountError,
+    TrackingUnavailableError,
+    RateLimitExceededError,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # Core
+    "CostTracker",
+    "Budget",
+    "BudgetAction",
+    "RateLimit",
+    "Span",
+    # Models
+    "CostRecord",
+    "CostReport",
+    "HealthStatus",
+    # Exceptions
+    "LLMCostGuardError",
+    "BudgetExceededError",
+    "PricingNotFoundError",
+    "TokenCountError",
+    "TrackingUnavailableError",
+    "RateLimitExceededError",
+]

llm_cost_guard/backends/__init__.py

@@ -0,0 +1,52 @@
+"""
+Storage backends for LLM Cost Guard.
+"""
+
+from llm_cost_guard.backends.base import Backend
+from llm_cost_guard.backends.memory import MemoryBackend
+
+__all__ = ["Backend", "MemoryBackend", "get_backend"]
+
+
+def get_backend(backend_url: str, **kwargs) -> Backend:
+    """
+    Create a backend instance from a URL.
+
+    Supported formats:
+    - "memory" or "memory://" - In-memory storage
+    - "sqlite:///path/to/db.sqlite" - SQLite database
+    - "postgresql://user:pass@host/db" - PostgreSQL database
+    - "redis://host:port/db" - Redis
+    - "dynamodb://table-name" - DynamoDB
+
+    Args:
+        backend_url: Backend connection URL
+        **kwargs: Additional backend-specific options
+
+    Returns:
+        Backend instance
+    """
+    if backend_url in ("memory", "memory://"):
+        return MemoryBackend(**kwargs)
+
+    if backend_url.startswith("sqlite:"):
+        from llm_cost_guard.backends.sqlite import SQLiteBackend
+
+        return SQLiteBackend(backend_url, **kwargs)
+
+    if backend_url.startswith("postgresql://") or backend_url.startswith("postgres://"):
+        from llm_cost_guard.backends.postgres import PostgresBackend
+
+        return PostgresBackend(backend_url, **kwargs)
+
+    if backend_url.startswith("redis://"):
+        from llm_cost_guard.backends.redis import RedisBackend
+
+        return RedisBackend(backend_url, **kwargs)
+
+    if backend_url.startswith("dynamodb://"):
+        from llm_cost_guard.backends.dynamodb import DynamoDBBackend
+
+        return DynamoDBBackend(backend_url, **kwargs)
+
+    raise ValueError(f"Unsupported backend URL: {backend_url}")

llm_cost_guard/backends/base.py

@@ -0,0 +1,121 @@
+"""
+Base backend interface for LLM Cost Guard.
+"""
+
+from abc import ABC, abstractmethod
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from llm_cost_guard.models import CostRecord, CostReport
+
+
+class Backend(ABC):
+    """Abstract base class for storage backends."""
+
+    @abstractmethod
+    def save_record(self, record: CostRecord) -> None:
+        """Save a cost record."""
+        pass
+
+    @abstractmethod
+    def save_records(self, records: List[CostRecord]) -> None:
+        """Save multiple cost records."""
+        pass
+
+    @abstractmethod
+    def get_records(
+        self,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        tags: Optional[Dict[str, str]] = None,
+        limit: Optional[int] = None,
+        offset: int = 0,
+    ) -> List[CostRecord]:
+        """
+        Retrieve cost records with optional filters.
+
+        Args:
+            start_date: Filter records after this date
+            end_date: Filter records before this date
+            tags: Filter by tag key-value pairs
+            limit: Maximum number of records to return
+            offset: Number of records to skip
+
+        Returns:
+            List of matching CostRecord objects
+        """
+        pass
+
+    @abstractmethod
+    def get_total_cost(
+        self,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        tags: Optional[Dict[str, str]] = None,
+    ) -> float:
+        """Get total cost for the given filters."""
+        pass
+
+    @abstractmethod
+    def get_aggregated_costs(
+        self,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        tags: Optional[Dict[str, str]] = None,
+        group_by: Optional[List[str]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Get aggregated costs grouped by specified fields.
+
+        Args:
+            start_date: Filter records after this date
+            end_date: Filter records before this date
+            tags: Filter by tag key-value pairs
+            group_by: Fields to group by (e.g., ["provider", "model", "team"])
+
+        Returns:
+            Dictionary with grouped cost data
+        """
+        pass
+
+    @abstractmethod
+    def get_report(
+        self,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        tags: Optional[Dict[str, str]] = None,
+        group_by: Optional[List[str]] = None,
+    ) -> CostReport:
+        """Generate a cost report."""
+        pass
+
+    @abstractmethod
+    def delete_records(
+        self,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        tags: Optional[Dict[str, str]] = None,
+    ) -> int:
+        """
+        Delete records matching the filters.
+
+        Returns:
+            Number of records deleted
+        """
+        pass
+
+    @abstractmethod
+    def health_check(self) -> bool:
+        """Check if the backend is healthy and connected."""
+        pass
+
+    @abstractmethod
+    def close(self) -> None:
+        """Close the backend connection."""
+        pass
+
+    def __enter__(self) -> "Backend":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()

llm_cost_guard/backends/memory.py

@@ -0,0 +1,265 @@
+"""
+In-memory storage backend for LLM Cost Guard.
+"""
+
+from collections import defaultdict
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+import threading
+
+from llm_cost_guard.backends.base import Backend
+from llm_cost_guard.models import CostRecord, CostReport
+
+
+class MemoryBackend(Backend):
+    """Thread-safe in-memory storage backend."""
+
+    def __init__(self, max_records: int = 100000, **kwargs):
+        """
+        Initialize the memory backend.
+
+        Args:
+            max_records: Maximum number of records to keep in memory
+        """
+        self._records: List[CostRecord] = []
+        self._max_records = max_records
+        self._lock = threading.RLock()
+
+    def save_record(self, record: CostRecord) -> None:
+        """Save a cost record."""
+        with self._lock:
+            self._records.append(record)
+            # Evict old records if we exceed the limit
+            if len(self._records) > self._max_records:
+                # Remove oldest 10% of records
+                evict_count = self._max_records // 10
+                self._records = self._records[evict_count:]
+
+    def save_records(self, records: List[CostRecord]) -> None:
+        """Save multiple cost records."""
+        with self._lock:
+            self._records.extend(records)
+            # Evict old records if we exceed the limit
+            if len(self._records) > self._max_records:
+                evict_count = len(self._records) - self._max_records
+                self._records = self._records[evict_count:]
+
+    def _matches_filters(
+        self,
+        record: CostRecord,
+        start_date: Optional[datetime],
+        end_date: Optional[datetime],
+        tags: Optional[Dict[str, str]],
+    ) -> bool:
+        """Check if a record matches the given filters."""
+        if start_date and record.timestamp < start_date:
+            return False
+        if end_date and record.timestamp > end_date:
+            return False
+        if tags:
+            for key, value in tags.items():
+                if record.tags.get(key) != value:
+                    return False
+        return True
+
+    def get_records(
+        self,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        tags: Optional[Dict[str, str]] = None,
+        limit: Optional[int] = None,
+        offset: int = 0,
+    ) -> List[CostRecord]:
+        """Retrieve cost records with optional filters."""
+        with self._lock:
+            filtered = [
+                r for r in self._records if self._matches_filters(r, start_date, end_date, tags)
+            ]
+
+            # Sort by timestamp descending (most recent first)
+            filtered.sort(key=lambda r: r.timestamp, reverse=True)
+
+            # Apply offset and limit
+            if offset:
+                filtered = filtered[offset:]
+            if limit:
+                filtered = filtered[:limit]
+
+            return filtered
+
+    def get_total_cost(
+        self,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        tags: Optional[Dict[str, str]] = None,
+    ) -> float:
+        """Get total cost for the given filters."""
+        with self._lock:
+            total = 0.0
+            for record in self._records:
+                if self._matches_filters(record, start_date, end_date, tags):
+                    total += record.total_cost
+            return total
+
+    def get_aggregated_costs(
+        self,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        tags: Optional[Dict[str, str]] = None,
+        group_by: Optional[List[str]] = None,
+    ) -> Dict[str, Any]:
+        """Get aggregated costs grouped by specified fields."""
+        with self._lock:
+            if not group_by:
+                # Return overall totals
+                total_cost = 0.0
+                total_calls = 0
+                total_input_tokens = 0
+                total_output_tokens = 0
+
+                for record in self._records:
+                    if self._matches_filters(record, start_date, end_date, tags):
+                        total_cost += record.total_cost
+                        total_calls += 1
+                        total_input_tokens += record.input_tokens
+                        total_output_tokens += record.output_tokens
+
+                return {
+                    "total_cost": total_cost,
+                    "total_calls": total_calls,
+                    "total_input_tokens": total_input_tokens,
+                    "total_output_tokens": total_output_tokens,
+                }
+
+            # Group by specified fields
+            groups: Dict[tuple, Dict[str, Any]] = defaultdict(
+                lambda: {
+                    "cost": 0.0,
+                    "calls": 0,
+                    "input_tokens": 0,
+                    "output_tokens": 0,
+                }
+            )
+
+            for record in self._records:
+                if not self._matches_filters(record, start_date, end_date, tags):
+                    continue
+
+                # Build group key
+                key_parts = []
+                for field in group_by:
+                    if field == "provider":
+                        key_parts.append(record.provider)
+                    elif field == "model":
+                        key_parts.append(record.model)
+                    elif field in record.tags:
+                        key_parts.append(record.tags[field])
+                    else:
+                        key_parts.append("unknown")
+
+                key = tuple(key_parts)
+                groups[key]["cost"] += record.total_cost
+                groups[key]["calls"] += 1
+                groups[key]["input_tokens"] += record.input_tokens
+                groups[key]["output_tokens"] += record.output_tokens
+
+            # Convert to list format
+            result = []
+            for key, data in groups.items():
+                row = dict(zip(group_by, key))
+                row.update(data)
+                result.append(row)
+
+            # Sort by cost descending
+            result.sort(key=lambda x: x["cost"], reverse=True)
+
+            return {"groups": result, "group_by": group_by}
+
+    def get_report(
+        self,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        tags: Optional[Dict[str, str]] = None,
+        group_by: Optional[List[str]] = None,
+    ) -> CostReport:
+        """Generate a cost report."""
+        with self._lock:
+            records = self.get_records(start_date, end_date, tags)
+
+            total_cost = 0.0
+            total_input_tokens = 0
+            total_output_tokens = 0
+            successful_calls = 0
+            failed_calls = 0
+            cache_hits = 0
+            cache_savings = 0.0
+
+            for record in records:
+                total_cost += record.total_cost
+                total_input_tokens += record.input_tokens
+                total_output_tokens += record.output_tokens
+
+                if record.success:
+                    successful_calls += 1
+                else:
+                    failed_calls += 1
+
+                if record.cached:
+                    cache_hits += 1
+                    cache_savings += record.cache_savings
+
+            grouped_data = {}
+            if group_by:
+                agg = self.get_aggregated_costs(start_date, end_date, tags, group_by)
+                grouped_data = agg.get("groups", [])
+
+            return CostReport(
+                start_date=start_date,
+                end_date=end_date,
+                total_cost=total_cost,
+                total_input_tokens=total_input_tokens,
+                total_output_tokens=total_output_tokens,
+                total_calls=len(records),
+                successful_calls=successful_calls,
+                failed_calls=failed_calls,
+                cache_hits=cache_hits,
+                cache_savings=cache_savings,
+                effective_cost=total_cost - cache_savings,
+                records=records,
+                grouped_data={"groups": grouped_data} if grouped_data else {},
+            )
+
+    def delete_records(
+        self,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        tags: Optional[Dict[str, str]] = None,
+    ) -> int:
+        """Delete records matching the filters."""
+        with self._lock:
+            initial_count = len(self._records)
+            self._records = [
+                r
+                for r in self._records
+                if not self._matches_filters(r, start_date, end_date, tags)
+            ]
+            return initial_count - len(self._records)
+
+    def health_check(self) -> bool:
+        """Check if the backend is healthy."""
+        return True
+
+    def close(self) -> None:
+        """Close the backend (no-op for memory backend)."""
+        pass
+
+    def clear(self) -> None:
+        """Clear all records."""
+        with self._lock:
+            self._records.clear()
+
+    @property
+    def record_count(self) -> int:
+        """Get the current number of records."""
+        with self._lock:
+            return len(self._records)