llm-cost-guard 0.1.2__py3-none-any.whl → 0.2.0__py3-none-any.whl

llm_cost_guard/__init__.py

@@ -15,8 +15,16 @@ from llm_cost_guard.exceptions import (
     TrackingUnavailableError,
     RateLimitExceededError,
 )
+from llm_cost_guard.audit import (
+    AuditLogger,
+    AuditBackend,
+    AuditEvent,
+    AuditEventType,
+    LoggingAuditBackend,
+    FileAuditBackend,
+)
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 
 __all__ = [
     # Core
@@ -29,6 +37,13 @@ __all__ = [
     "CostRecord",
     "CostReport",
     "HealthStatus",
+    # Audit
+    "AuditLogger",
+    "AuditBackend",
+    "AuditEvent",
+    "AuditEventType",
+    "LoggingAuditBackend",
+    "FileAuditBackend",
     # Exceptions
     "LLMCostGuardError",
     "BudgetExceededError",

llm_cost_guard/audit.py

@@ -0,0 +1,480 @@
+"""
+Audit logging for LLM Cost Guard.
+
+Provides compliance-ready audit trails for:
+- Budget changes (created, modified, deleted)
+- Budget exceeded events
+- Configuration changes
+- Rate limit events
+"""
+
+import json
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field, asdict
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Optional, Callable
+
+logger = logging.getLogger(__name__)
+
+
+class AuditEventType(str, Enum):
+    """Types of audit events."""
+    
+    # Budget events
+    BUDGET_CREATED = "budget.created"
+    BUDGET_MODIFIED = "budget.modified"
+    BUDGET_DELETED = "budget.deleted"
+    BUDGET_WARNING = "budget.warning"
+    BUDGET_EXCEEDED = "budget.exceeded"
+    BUDGET_RESET = "budget.reset"
+    
+    # Rate limit events
+    RATE_LIMIT_EXCEEDED = "rate_limit.exceeded"
+    
+    # Configuration events
+    CONFIG_CHANGED = "config.changed"
+    BACKEND_CHANGED = "backend.changed"
+    PRICING_UPDATED = "pricing.updated"
+    
+    # Tracking events
+    TRACKING_FAILURE = "tracking.failure"
+    FALLBACK_ACTIVATED = "fallback.activated"
+    
+    # Cost events
+    COST_RECORDED = "cost.recorded"
+    COST_ANOMALY = "cost.anomaly"
+
+
+@dataclass
+class AuditEvent:
+    """Represents an audit event."""
+    
+    event_type: AuditEventType
+    timestamp: datetime = field(default_factory=datetime.now)
+    actor: Optional[str] = None  # Who initiated the action (user/system)
+    resource: Optional[str] = None  # What was affected (budget name, etc.)
+    details: Dict[str, Any] = field(default_factory=dict)
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "event_type": self.event_type.value,
+            "timestamp": self.timestamp.isoformat(),
+            "actor": self.actor,
+            "resource": self.resource,
+            "details": self.details,
+            "metadata": self.metadata,
+        }
+    
+    def to_json(self) -> str:
+        """Convert to JSON string."""
+        return json.dumps(self.to_dict())
+
+
+class AuditBackend(ABC):
+    """Abstract base class for audit backends."""
+    
+    @abstractmethod
+    def log(self, event: AuditEvent) -> None:
+        """Log an audit event."""
+        pass
+    
+    @abstractmethod
+    def query(
+        self,
+        event_type: Optional[AuditEventType] = None,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        resource: Optional[str] = None,
+        limit: int = 100,
+    ) -> List[AuditEvent]:
+        """Query audit events."""
+        pass
+
+
+class LoggingAuditBackend(AuditBackend):
+    """
+    Audit backend that logs to Python logging.
+    
+    Suitable for development and when audit logs should go to 
+    existing log aggregation systems (CloudWatch, DataDog, etc.).
+    """
+    
+    def __init__(
+        self,
+        logger_name: str = "llm_cost_guard.audit",
+        log_level: int = logging.INFO,
+    ):
+        self._logger = logging.getLogger(logger_name)
+        self._log_level = log_level
+        self._events: List[AuditEvent] = []  # In-memory for query support
+        self._max_events = 10000  # Limit in-memory storage
+    
+    def log(self, event: AuditEvent) -> None:
+        """Log an audit event."""
+        # Log to Python logger
+        self._logger.log(
+            self._log_level,
+            f"AUDIT: {event.event_type.value} | resource={event.resource} | {event.details}"
+        )
+        
+        # Store in memory for queries
+        self._events.append(event)
+        
+        # Trim if too many events
+        if len(self._events) > self._max_events:
+            self._events = self._events[-self._max_events:]
+    
+    def query(
+        self,
+        event_type: Optional[AuditEventType] = None,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        resource: Optional[str] = None,
+        limit: int = 100,
+    ) -> List[AuditEvent]:
+        """Query audit events."""
+        results = []
+        
+        for event in reversed(self._events):  # Most recent first
+            if event_type and event.event_type != event_type:
+                continue
+            if start_date and event.timestamp < start_date:
+                continue
+            if end_date and event.timestamp > end_date:
+                continue
+            if resource and event.resource != resource:
+                continue
+            
+            results.append(event)
+            
+            if len(results) >= limit:
+                break
+        
+        return results
+
+
+class FileAuditBackend(AuditBackend):
+    """
+    Audit backend that writes to a file (JSON Lines format).
+    
+    Suitable for compliance requirements where audit logs need
+    to be stored in a specific location.
+    """
+    
+    def __init__(self, file_path: str):
+        self._file_path = file_path
+    
+    def log(self, event: AuditEvent) -> None:
+        """Log an audit event."""
+        with open(self._file_path, "a") as f:
+            f.write(event.to_json() + "\n")
+    
+    def query(
+        self,
+        event_type: Optional[AuditEventType] = None,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        resource: Optional[str] = None,
+        limit: int = 100,
+    ) -> List[AuditEvent]:
+        """Query audit events from file."""
+        results = []
+        
+        try:
+            with open(self._file_path, "r") as f:
+                lines = f.readlines()
+        except FileNotFoundError:
+            return []
+        
+        for line in reversed(lines):  # Most recent first
+            if not line.strip():
+                continue
+            
+            data = json.loads(line)
+            event = AuditEvent(
+                event_type=AuditEventType(data["event_type"]),
+                timestamp=datetime.fromisoformat(data["timestamp"]),
+                actor=data.get("actor"),
+                resource=data.get("resource"),
+                details=data.get("details", {}),
+                metadata=data.get("metadata", {}),
+            )
+            
+            if event_type and event.event_type != event_type:
+                continue
+            if start_date and event.timestamp < start_date:
+                continue
+            if end_date and event.timestamp > end_date:
+                continue
+            if resource and event.resource != resource:
+                continue
+            
+            results.append(event)
+            
+            if len(results) >= limit:
+                break
+        
+        return results
+
+
+class CompositeAuditBackend(AuditBackend):
+    """
+    Audit backend that writes to multiple backends.
+    
+    Useful for sending audit logs to both local storage and remote systems.
+    """
+    
+    def __init__(self, backends: List[AuditBackend]):
+        self._backends = backends
+    
+    def log(self, event: AuditEvent) -> None:
+        """Log an audit event to all backends."""
+        for backend in self._backends:
+            try:
+                backend.log(event)
+            except Exception as e:
+                logger.error(f"Audit backend failed: {e}")
+    
+    def query(
+        self,
+        event_type: Optional[AuditEventType] = None,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        resource: Optional[str] = None,
+        limit: int = 100,
+    ) -> List[AuditEvent]:
+        """Query from first backend that succeeds."""
+        for backend in self._backends:
+            try:
+                return backend.query(event_type, start_date, end_date, resource, limit)
+            except Exception as e:
+                logger.warning(f"Audit query failed on backend: {e}")
+        return []
+
+
+class AuditLogger:
+    """
+    Main audit logger for LLM Cost Guard.
+    
+    Usage:
+        audit = AuditLogger(backend=FileAuditBackend("audit.log"))
+        
+        audit.log_budget_created(budget)
+        audit.log_budget_exceeded(budget, current_spending)
+    """
+    
+    def __init__(
+        self,
+        backend: Optional[AuditBackend] = None,
+        enabled: bool = True,
+        actor: Optional[str] = None,
+    ):
+        self._backend = backend or LoggingAuditBackend()
+        self._enabled = enabled
+        self._default_actor = actor or "system"
+        self._event_callbacks: Dict[AuditEventType, List[Callable[[AuditEvent], None]]] = {}
+    
+    def _log(self, event: AuditEvent) -> None:
+        """Internal logging method."""
+        if not self._enabled:
+            return
+        
+        if event.actor is None:
+            event.actor = self._default_actor
+        
+        self._backend.log(event)
+        
+        # Trigger callbacks
+        callbacks = self._event_callbacks.get(event.event_type, [])
+        for callback in callbacks:
+            try:
+                callback(event)
+            except Exception as e:
+                logger.error(f"Audit callback failed: {e}")
+    
+    def on_event(
+        self,
+        event_type: AuditEventType,
+        callback: Callable[[AuditEvent], None],
+    ) -> None:
+        """Register a callback for an event type."""
+        if event_type not in self._event_callbacks:
+            self._event_callbacks[event_type] = []
+        self._event_callbacks[event_type].append(callback)
+    
+    # Budget events
+    def log_budget_created(
+        self,
+        budget_name: str,
+        limit: float,
+        period: str,
+        action: str,
+        actor: Optional[str] = None,
+    ) -> None:
+        """Log budget creation."""
+        self._log(AuditEvent(
+            event_type=AuditEventType.BUDGET_CREATED,
+            actor=actor,
+            resource=budget_name,
+            details={
+                "limit": limit,
+                "period": period,
+                "action": action,
+            },
+        ))
+    
+    def log_budget_modified(
+        self,
+        budget_name: str,
+        old_values: Dict[str, Any],
+        new_values: Dict[str, Any],
+        actor: Optional[str] = None,
+    ) -> None:
+        """Log budget modification."""
+        self._log(AuditEvent(
+            event_type=AuditEventType.BUDGET_MODIFIED,
+            actor=actor,
+            resource=budget_name,
+            details={
+                "old_values": old_values,
+                "new_values": new_values,
+            },
+        ))
+    
+    def log_budget_deleted(
+        self,
+        budget_name: str,
+        actor: Optional[str] = None,
+    ) -> None:
+        """Log budget deletion."""
+        self._log(AuditEvent(
+            event_type=AuditEventType.BUDGET_DELETED,
+            actor=actor,
+            resource=budget_name,
+        ))
+    
+    def log_budget_warning(
+        self,
+        budget_name: str,
+        current_spending: float,
+        limit: float,
+        utilization: float,
+    ) -> None:
+        """Log budget warning."""
+        self._log(AuditEvent(
+            event_type=AuditEventType.BUDGET_WARNING,
+            resource=budget_name,
+            details={
+                "current_spending": current_spending,
+                "limit": limit,
+                "utilization_percent": utilization * 100,
+            },
+        ))
+    
+    def log_budget_exceeded(
+        self,
+        budget_name: str,
+        current_spending: float,
+        limit: float,
+        action_taken: str,
+    ) -> None:
+        """Log budget exceeded event."""
+        self._log(AuditEvent(
+            event_type=AuditEventType.BUDGET_EXCEEDED,
+            resource=budget_name,
+            details={
+                "current_spending": current_spending,
+                "limit": limit,
+                "action_taken": action_taken,
+            },
+        ))
+    
+    def log_budget_reset(
+        self,
+        budget_name: str,
+        reason: str = "period_end",
+        actor: Optional[str] = None,
+    ) -> None:
+        """Log budget reset."""
+        self._log(AuditEvent(
+            event_type=AuditEventType.BUDGET_RESET,
+            actor=actor,
+            resource=budget_name,
+            details={"reason": reason},
+        ))
+    
+    # Rate limit events
+    def log_rate_limit_exceeded(
+        self,
+        limit_name: str,
+        current: int,
+        limit: int,
+        retry_after: float,
+    ) -> None:
+        """Log rate limit exceeded event."""
+        self._log(AuditEvent(
+            event_type=AuditEventType.RATE_LIMIT_EXCEEDED,
+            resource=limit_name,
+            details={
+                "current": current,
+                "limit": limit,
+                "retry_after_seconds": retry_after,
+            },
+        ))
+    
+    # Tracking events
+    def log_tracking_failure(
+        self,
+        error: str,
+        backend: str,
+        action_taken: str,
+    ) -> None:
+        """Log tracking failure."""
+        self._log(AuditEvent(
+            event_type=AuditEventType.TRACKING_FAILURE,
+            resource=backend,
+            details={
+                "error": error,
+                "action_taken": action_taken,
+            },
+        ))
+    
+    def log_fallback_activated(
+        self,
+        original_backend: str,
+        fallback_backend: str,
+        reason: str,
+    ) -> None:
+        """Log fallback activation."""
+        self._log(AuditEvent(
+            event_type=AuditEventType.FALLBACK_ACTIVATED,
+            details={
+                "original_backend": original_backend,
+                "fallback_backend": fallback_backend,
+                "reason": reason,
+            },
+        ))
+    
+    # Query methods
+    def query(
+        self,
+        event_type: Optional[AuditEventType] = None,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        resource: Optional[str] = None,
+        limit: int = 100,
+    ) -> List[AuditEvent]:
+        """Query audit events."""
+        return self._backend.query(event_type, start_date, end_date, resource, limit)
+    
+    def get_budget_history(
+        self,
+        budget_name: str,
+        limit: int = 100,
+    ) -> List[AuditEvent]:
+        """Get audit history for a specific budget."""
+        return self._backend.query(resource=budget_name, limit=limit)

llm_cost_guard/backends/__init__.py

@@ -40,7 +40,7 @@ def get_backend(backend_url: str, **kwargs) -> Backend:
         return PostgresBackend(backend_url, **kwargs)
 
     if backend_url.startswith("redis://"):
-        from llm_cost_guard.backends.redis import RedisBackend
+        from llm_cost_guard.backends.redis_backend import RedisBackend
 
         return RedisBackend(backend_url, **kwargs)