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

llm_cost_guard/providers/base.py

@@ -0,0 +1,72 @@
+"""
+Base provider interface for LLM Cost Guard.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional
+
+from llm_cost_guard.models import UsageData
+
+
+class Provider(ABC):
+    """Abstract base class for LLM providers."""
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Get the provider name."""
+        pass
+
+    @abstractmethod
+    def extract_usage(self, response: Any) -> UsageData:
+        """
+        Extract token usage from an API response.
+
+        Args:
+            response: The raw API response object
+
+        Returns:
+            UsageData with token counts
+        """
+        pass
+
+    @abstractmethod
+    def extract_model(self, response: Any) -> str:
+        """
+        Extract the model name from an API response.
+
+        Args:
+            response: The raw API response object
+
+        Returns:
+            Model name string
+        """
+        pass
+
+    def extract_cached_tokens(self, response: Any) -> int:
+        """
+        Extract cached token count from an API response.
+
+        Args:
+            response: The raw API response object
+
+        Returns:
+            Number of cached tokens (0 if not applicable)
+        """
+        return 0
+
+    def supports_streaming(self) -> bool:
+        """Check if this provider supports streaming."""
+        return True
+
+    def normalize_model_name(self, model: str) -> str:
+        """
+        Normalize model name for consistent pricing lookup.
+
+        Args:
+            model: Raw model name from API
+
+        Returns:
+            Normalized model name
+        """
+        return model

llm_cost_guard/providers/bedrock.py

@@ -0,0 +1,135 @@
+"""
+AWS Bedrock provider for LLM Cost Guard.
+"""
+
+import json
+from typing import Any
+
+from llm_cost_guard.models import UsageData
+from llm_cost_guard.providers.base import Provider
+
+
+class BedrockProvider(Provider):
+    """AWS Bedrock API provider."""
+
+    @property
+    def name(self) -> str:
+        return "bedrock"
+
+    def extract_usage(self, response: Any) -> UsageData:
+        """Extract token usage from a Bedrock API response."""
+        usage = UsageData()
+
+        # Handle dictionary response (from boto3)
+        if isinstance(response, dict):
+            # Check for standard usage field
+            if "usage" in response:
+                usage_data = response["usage"]
+                usage.input_tokens = usage_data.get("inputTokens", 0)
+                usage.output_tokens = usage_data.get("outputTokens", 0)
+                usage.total_tokens = usage_data.get("totalTokens", 0)
+                return usage
+
+            # Check for Bedrock response metadata
+            metadata = response.get("ResponseMetadata", {})
+            headers = metadata.get("HTTPHeaders", {})
+
+            # Bedrock includes token counts in headers for some models
+            if "x-amzn-bedrock-input-token-count" in headers:
+                usage.input_tokens = int(headers["x-amzn-bedrock-input-token-count"])
+            if "x-amzn-bedrock-output-token-count" in headers:
+                usage.output_tokens = int(headers["x-amzn-bedrock-output-token-count"])
+
+            # Also check the body for Claude-style responses
+            body = response.get("body", {})
+            if isinstance(body, bytes):
+                try:
+                    body = json.loads(body.decode("utf-8"))
+                except (json.JSONDecodeError, UnicodeDecodeError):
+                    body = {}
+
+            if isinstance(body, dict) and "usage" in body:
+                body_usage = body["usage"]
+                usage.input_tokens = body_usage.get(
+                    "input_tokens", body_usage.get("inputTokens", usage.input_tokens)
+                )
+                usage.output_tokens = body_usage.get(
+                    "output_tokens", body_usage.get("outputTokens", usage.output_tokens)
+                )
+
+            usage.total_tokens = usage.input_tokens + usage.output_tokens
+
+        return usage
+
+    def extract_model(self, response: Any) -> str:
+        """Extract the model name from a Bedrock API response."""
+        if isinstance(response, dict):
+            # Check for model ID in various locations
+            if "modelId" in response:
+                return response["modelId"]
+
+            # Check response metadata
+            metadata = response.get("ResponseMetadata", {})
+            if "modelId" in metadata:
+                return metadata["modelId"]
+
+        return "unknown"
+
+    def normalize_model_name(self, model: str) -> str:
+        """Normalize Bedrock model name for pricing lookup."""
+        # Bedrock model IDs include version info that should be preserved
+        # e.g., "anthropic.claude-3-sonnet-20240229-v1:0"
+        return model
+
+
+class BedrockStreamingHandler:
+    """Handler for streaming Bedrock responses."""
+
+    def __init__(self, model_id: str = "unknown"):
+        self.input_tokens = 0
+        self.output_tokens = 0
+        self.model = model_id
+        self._chunks: list = []
+
+    def handle_chunk(self, chunk: Any) -> None:
+        """Process a streaming chunk."""
+        if isinstance(chunk, dict):
+            # Handle Bedrock's streaming chunk format
+            if "chunk" in chunk:
+                chunk_data = chunk["chunk"]
+                if "bytes" in chunk_data:
+                    try:
+                        parsed = json.loads(chunk_data["bytes"].decode("utf-8"))
+                        self._process_parsed_chunk(parsed)
+                    except (json.JSONDecodeError, UnicodeDecodeError, AttributeError):
+                        pass
+
+            # Some responses have direct usage info
+            if "usage" in chunk:
+                usage = chunk["usage"]
+                self.input_tokens = usage.get("inputTokens", self.input_tokens)
+                self.output_tokens = usage.get("outputTokens", self.output_tokens)
+
+        self._chunks.append(chunk)
+
+    def _process_parsed_chunk(self, parsed: dict) -> None:
+        """Process a parsed JSON chunk."""
+        # Claude on Bedrock format
+        if "usage" in parsed:
+            usage = parsed["usage"]
+            self.input_tokens = usage.get("input_tokens", self.input_tokens)
+            self.output_tokens = usage.get("output_tokens", self.output_tokens)
+
+        # Llama/Titan format
+        if "amazon-bedrock-invocationMetrics" in parsed:
+            metrics = parsed["amazon-bedrock-invocationMetrics"]
+            self.input_tokens = metrics.get("inputTokenCount", self.input_tokens)
+            self.output_tokens = metrics.get("outputTokenCount", self.output_tokens)
+
+    def get_usage(self) -> UsageData:
+        """Get final usage data."""
+        return UsageData(
+            input_tokens=self.input_tokens,
+            output_tokens=self.output_tokens,
+            total_tokens=self.input_tokens + self.output_tokens,
+        )

llm_cost_guard/providers/openai.py

@@ -0,0 +1,110 @@
+"""
+OpenAI provider for LLM Cost Guard.
+"""
+
+from typing import Any, Optional
+
+from llm_cost_guard.models import UsageData
+from llm_cost_guard.providers.base import Provider
+
+
+class OpenAIProvider(Provider):
+    """OpenAI API provider."""
+
+    @property
+    def name(self) -> str:
+        return "openai"
+
+    def extract_usage(self, response: Any) -> UsageData:
+        """Extract token usage from an OpenAI API response."""
+        usage = UsageData()
+
+        # Handle dictionary response
+        if isinstance(response, dict):
+            usage_data = response.get("usage", {})
+            usage.input_tokens = usage_data.get("prompt_tokens", 0)
+            usage.output_tokens = usage_data.get("completion_tokens", 0)
+            usage.total_tokens = usage_data.get("total_tokens", 0)
+
+            # Check for cached tokens (prompt caching)
+            prompt_tokens_details = usage_data.get("prompt_tokens_details", {})
+            if prompt_tokens_details:
+                usage.cached_tokens = prompt_tokens_details.get("cached_tokens", 0)
+
+            return usage
+
+        # Handle OpenAI client response object
+        if hasattr(response, "usage") and response.usage is not None:
+            usage.input_tokens = getattr(response.usage, "prompt_tokens", 0) or 0
+            usage.output_tokens = getattr(response.usage, "completion_tokens", 0) or 0
+            usage.total_tokens = getattr(response.usage, "total_tokens", 0) or 0
+
+            # Check for prompt caching details
+            if hasattr(response.usage, "prompt_tokens_details"):
+                details = response.usage.prompt_tokens_details
+                if details and hasattr(details, "cached_tokens"):
+                    usage.cached_tokens = details.cached_tokens or 0
+
+        return usage
+
+    def extract_model(self, response: Any) -> str:
+        """Extract the model name from an OpenAI API response."""
+        if isinstance(response, dict):
+            return response.get("model", "unknown")
+
+        if hasattr(response, "model"):
+            return response.model or "unknown"
+
+        return "unknown"
+
+    def extract_cached_tokens(self, response: Any) -> int:
+        """Extract cached token count from an OpenAI API response."""
+        usage = self.extract_usage(response)
+        return usage.cached_tokens
+
+    def normalize_model_name(self, model: str) -> str:
+        """Normalize OpenAI model name."""
+        # Remove date suffixes for pricing lookup
+        # e.g., "gpt-4-0613" -> "gpt-4"
+        parts = model.split("-")
+        if len(parts) >= 3 and parts[-1].isdigit() and len(parts[-1]) >= 4:
+            return "-".join(parts[:-1])
+        return model
+
+
+class OpenAIStreamingHandler:
+    """Handler for streaming OpenAI responses."""
+
+    def __init__(self):
+        self.input_tokens = 0
+        self.output_tokens = 0
+        self.model = "unknown"
+        self._chunk_count = 0
+
+    def handle_chunk(self, chunk: Any) -> None:
+        """Process a streaming chunk."""
+        self._chunk_count += 1
+
+        # Extract model from first chunk
+        if self._chunk_count == 1:
+            if isinstance(chunk, dict):
+                self.model = chunk.get("model", self.model)
+            elif hasattr(chunk, "model"):
+                self.model = chunk.model or self.model
+
+        # Some providers include usage in the final chunk
+        if isinstance(chunk, dict):
+            if "usage" in chunk and chunk["usage"]:
+                self.input_tokens = chunk["usage"].get("prompt_tokens", 0)
+                self.output_tokens = chunk["usage"].get("completion_tokens", 0)
+        elif hasattr(chunk, "usage") and chunk.usage:
+            self.input_tokens = getattr(chunk.usage, "prompt_tokens", 0) or 0
+            self.output_tokens = getattr(chunk.usage, "completion_tokens", 0) or 0
+
+    def get_usage(self) -> UsageData:
+        """Get final usage data."""
+        return UsageData(
+            input_tokens=self.input_tokens,
+            output_tokens=self.output_tokens,
+            total_tokens=self.input_tokens + self.output_tokens,
+        )

llm_cost_guard/rate_limit.py

@@ -0,0 +1,233 @@
+"""
+Rate limiting for LLM Cost Guard.
+"""
+
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from typing import Dict, List, Literal, Optional, Tuple
+import threading
+import time
+
+RateLimitPeriod = Literal["second", "minute", "hour"]
+RateLimitScope = Literal["global", "model", "provider"]
+
+
+@dataclass
+class RateLimit:
+    """Rate limit configuration."""
+
+    name: str
+    limit: int
+    period: RateLimitPeriod = "minute"
+    scope: str = "global"  # "global", "model", "provider", or "tag:key_name"
+
+
+class SlidingWindowCounter:
+    """Sliding window rate limiter implementation."""
+
+    def __init__(self, window_size_seconds: float, limit: int):
+        self._window_size = window_size_seconds
+        self._limit = limit
+        self._requests: List[float] = []
+        self._lock = threading.Lock()
+
+    def _cleanup(self, now: float) -> None:
+        """Remove expired entries."""
+        cutoff = now - self._window_size
+        self._requests = [t for t in self._requests if t > cutoff]
+
+    def check(self) -> Tuple[bool, int, Optional[float]]:
+        """
+        Check if a request would be allowed.
+        Returns (allowed, current_count, retry_after_seconds).
+        """
+        now = time.time()
+        with self._lock:
+            self._cleanup(now)
+            current = len(self._requests)
+
+            if current >= self._limit:
+                # Calculate when the oldest request will expire
+                if self._requests:
+                    retry_after = self._requests[0] + self._window_size - now
+                    return False, current, max(0.0, retry_after)
+                return False, current, self._window_size
+
+            return True, current, None
+
+    def record(self) -> bool:
+        """
+        Record a request. Returns True if allowed, False if rate limited.
+        """
+        now = time.time()
+        with self._lock:
+            self._cleanup(now)
+
+            if len(self._requests) >= self._limit:
+                return False
+
+            self._requests.append(now)
+            return True
+
+    def get_count(self) -> int:
+        """Get current request count in the window."""
+        now = time.time()
+        with self._lock:
+            self._cleanup(now)
+            return len(self._requests)
+
+    def reset(self) -> None:
+        """Reset the counter."""
+        with self._lock:
+            self._requests = []
+
+
+class RateLimiter:
+    """Manages rate limiting across multiple limits and scopes."""
+
+    def __init__(self, rate_limits: Optional[List[RateLimit]] = None):
+        self._rate_limits = rate_limits or []
+        # Key: (limit_name, scope_value) -> SlidingWindowCounter
+        self._counters: Dict[Tuple[str, str], SlidingWindowCounter] = {}
+        self._lock = threading.Lock()
+
+    def _get_period_seconds(self, period: RateLimitPeriod) -> float:
+        """Get period in seconds."""
+        periods = {
+            "second": 1.0,
+            "minute": 60.0,
+            "hour": 3600.0,
+        }
+        return periods.get(period, 60.0)
+
+    def _get_scope_key(
+        self,
+        rate_limit: RateLimit,
+        model: Optional[str] = None,
+        provider: Optional[str] = None,
+        tags: Optional[Dict[str, str]] = None,
+    ) -> str:
+        """Get the scope key for a rate limit."""
+        scope = rate_limit.scope
+        tags = tags or {}
+
+        if scope == "global":
+            return "global"
+        elif scope == "model":
+            return f"model:{model or 'unknown'}"
+        elif scope == "provider":
+            return f"provider:{provider or 'unknown'}"
+        elif scope.startswith("tag:"):
+            tag_key = scope[4:]
+            tag_value = tags.get(tag_key, "unknown")
+            return f"tag:{tag_key}:{tag_value}"
+        return "global"
+
+    def _get_counter(
+        self,
+        rate_limit: RateLimit,
+        scope_key: str,
+    ) -> SlidingWindowCounter:
+        """Get or create a counter for a rate limit and scope."""
+        key = (rate_limit.name, scope_key)
+
+        with self._lock:
+            if key not in self._counters:
+                window_size = self._get_period_seconds(rate_limit.period)
+                self._counters[key] = SlidingWindowCounter(window_size, rate_limit.limit)
+            return self._counters[key]
+
+    def add_rate_limit(self, rate_limit: RateLimit) -> None:
+        """Add a new rate limit."""
+        self._rate_limits.append(rate_limit)
+
+    def remove_rate_limit(self, name: str) -> bool:
+        """Remove a rate limit by name."""
+        for i, rl in enumerate(self._rate_limits):
+            if rl.name == name:
+                del self._rate_limits[i]
+                # Remove associated counters
+                with self._lock:
+                    keys_to_remove = [k for k in self._counters if k[0] == name]
+                    for k in keys_to_remove:
+                        del self._counters[k]
+                return True
+        return False
+
+    def get_rate_limit(self, name: str) -> Optional[RateLimit]:
+        """Get a rate limit by name."""
+        for rl in self._rate_limits:
+            if rl.name == name:
+                return rl
+        return None
+
+    def check(
+        self,
+        model: Optional[str] = None,
+        provider: Optional[str] = None,
+        tags: Optional[Dict[str, str]] = None,
+    ) -> List[Tuple[RateLimit, int, Optional[float]]]:
+        """
+        Check all rate limits.
+        Returns list of (rate_limit, current_count, retry_after) for exceeded limits.
+        """
+        exceeded = []
+
+        for rate_limit in self._rate_limits:
+            scope_key = self._get_scope_key(rate_limit, model, provider, tags)
+            counter = self._get_counter(rate_limit, scope_key)
+            allowed, current, retry_after = counter.check()
+
+            if not allowed:
+                exceeded.append((rate_limit, current, retry_after))
+
+        return exceeded
+
+    def record(
+        self,
+        model: Optional[str] = None,
+        provider: Optional[str] = None,
+        tags: Optional[Dict[str, str]] = None,
+    ) -> List[Tuple[RateLimit, int, Optional[float]]]:
+        """
+        Record a request against all applicable rate limits.
+        Returns list of (rate_limit, current_count, retry_after) for limits that are now exceeded.
+        """
+        exceeded = []
+
+        for rate_limit in self._rate_limits:
+            scope_key = self._get_scope_key(rate_limit, model, provider, tags)
+            counter = self._get_counter(rate_limit, scope_key)
+
+            if not counter.record():
+                _, current, retry_after = counter.check()
+                exceeded.append((rate_limit, current, retry_after))
+
+        return exceeded
+
+    def get_remaining(
+        self,
+        name: str,
+        model: Optional[str] = None,
+        provider: Optional[str] = None,
+        tags: Optional[Dict[str, str]] = None,
+    ) -> int:
+        """Get remaining requests for a rate limit."""
+        rate_limit = self.get_rate_limit(name)
+        if rate_limit is None:
+            return 0
+
+        scope_key = self._get_scope_key(rate_limit, model, provider, tags)
+        counter = self._get_counter(rate_limit, scope_key)
+        return max(0, rate_limit.limit - counter.get_count())
+
+    def reset(self, name: Optional[str] = None) -> None:
+        """Reset counters for a specific rate limit or all rate limits."""
+        with self._lock:
+            if name:
+                keys_to_reset = [k for k in self._counters if k[0] == name]
+                for k in keys_to_reset:
+                    self._counters[k].reset()
+            else:
+                for counter in self._counters.values():
+                    counter.reset()

llm_cost_guard/span.py

@@ -0,0 +1,143 @@
+"""
+Hierarchical tracking spans for LLM Cost Guard.
+"""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Set, TYPE_CHECKING
+import contextvars
+import threading
+import uuid
+
+if TYPE_CHECKING:
+    from llm_cost_guard.models import CostRecord
+
+
+# Context variable to track the current span
+_current_span: contextvars.ContextVar[Optional["Span"]] = contextvars.ContextVar(
+    "current_span", default=None
+)
+
+
+@dataclass
+class Span:
+    """Hierarchical tracking span for grouping multiple LLM calls."""
+
+    name: str
+    span_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    parent_id: Optional[str] = None
+    start_time: Optional[datetime] = None
+    end_time: Optional[datetime] = None
+    tags: Dict[str, str] = field(default_factory=dict)
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    # Aggregated metrics
+    total_cost: float = 0.0
+    total_input_tokens: int = 0
+    total_output_tokens: int = 0
+    call_count: int = 0
+    models_used: Set[str] = field(default_factory=set)
+
+    # Hierarchy
+    children: List["Span"] = field(default_factory=list)
+    _records: List["CostRecord"] = field(default_factory=list)
+    _lock: threading.Lock = field(default_factory=threading.Lock)
+
+    # For context management
+    _token: Optional[contextvars.Token] = field(default=None, repr=False)
+    _previous_span: Optional["Span"] = field(default=None, repr=False)
+
+    def __enter__(self) -> "Span":
+        """Enter the span context."""
+        self.start_time = datetime.now()
+
+        # Save previous span and set this as current
+        self._previous_span = _current_span.get()
+        self._token = _current_span.set(self)
+
+        # If there's a parent span, register as child
+        if self._previous_span is not None:
+            self.parent_id = self._previous_span.span_id
+            self._previous_span._add_child(self)
+
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Exit the span context."""
+        self.end_time = datetime.now()
+
+        # Restore previous span
+        if self._token is not None:
+            _current_span.reset(self._token)
+
+        # Propagate costs to parent
+        if self._previous_span is not None:
+            self._previous_span._propagate_child_costs(self)
+
+    def _add_child(self, child: "Span") -> None:
+        """Add a child span."""
+        with self._lock:
+            self.children.append(child)
+
+    def _propagate_child_costs(self, child: "Span") -> None:
+        """Propagate child span costs to this span."""
+        with self._lock:
+            self.total_cost += child.total_cost
+            self.total_input_tokens += child.total_input_tokens
+            self.total_output_tokens += child.total_output_tokens
+            self.call_count += child.call_count
+            self.models_used.update(child.models_used)
+
+    def record_call(
+        self,
+        cost: float,
+        input_tokens: int,
+        output_tokens: int,
+        model: str,
+        record: Optional["CostRecord"] = None,
+    ) -> None:
+        """Record an LLM call in this span."""
+        with self._lock:
+            self.total_cost += cost
+            self.total_input_tokens += input_tokens
+            self.total_output_tokens += output_tokens
+            self.call_count += 1
+            self.models_used.add(model)
+            if record is not None:
+                self._records.append(record)
+
+    @property
+    def duration_ms(self) -> Optional[int]:
+        """Get the span duration in milliseconds."""
+        if self.start_time is None or self.end_time is None:
+            return None
+        delta = self.end_time - self.start_time
+        return int(delta.total_seconds() * 1000)
+
+    @property
+    def records(self) -> List["CostRecord"]:
+        """Get all records in this span."""
+        return list(self._records)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert span to dictionary."""
+        return {
+            "name": self.name,
+            "span_id": self.span_id,
+            "parent_id": self.parent_id,
+            "start_time": self.start_time.isoformat() if self.start_time else None,
+            "end_time": self.end_time.isoformat() if self.end_time else None,
+            "duration_ms": self.duration_ms,
+            "tags": self.tags,
+            "total_cost": self.total_cost,
+            "total_input_tokens": self.total_input_tokens,
+            "total_output_tokens": self.total_output_tokens,
+            "call_count": self.call_count,
+            "models_used": list(self.models_used),
+            "children": [child.to_dict() for child in self.children],
+        }
+
+
+def get_current_span() -> Optional[Span]:
+    """Get the current active span, if any."""
+    return _current_span.get()