stratifyai 0.1.0__py3-none-any.whl

stratifyai/cost_tracker.py

@@ -0,0 +1,257 @@
+"""Cost tracking module for LLM API calls."""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Dict, List, Optional
+from collections import defaultdict
+
+
+@dataclass
+class CostEntry:
+    """Individual cost entry for an LLM API call."""
+    
+    timestamp: datetime
+    provider: str
+    model: str
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    cost_usd: float
+    request_id: str
+    cached_tokens: int = 0
+    cache_creation_tokens: int = 0
+    cache_read_tokens: int = 0
+    group: Optional[str] = None
+
+
+class CostTracker:
+    """
+    Track and analyze costs across LLM API calls.
+    
+    Features:
+    - Call history with detailed metrics
+    - Grouping by provider, model, or custom tags
+    - Cost analytics and reporting
+    - Budget tracking and alerts
+    """
+    
+    def __init__(self):
+        """Initialize cost tracker."""
+        self._entries: List[CostEntry] = []
+        self._total_cost: float = 0.0
+        self._budget_limit: Optional[float] = None
+        self._alert_threshold: Optional[float] = None
+    
+    def add_entry(
+        self,
+        provider: str,
+        model: str,
+        prompt_tokens: int,
+        completion_tokens: int,
+        total_tokens: int,
+        cost_usd: float,
+        request_id: str,
+        cached_tokens: int = 0,
+        cache_creation_tokens: int = 0,
+        cache_read_tokens: int = 0,
+        group: Optional[str] = None,
+    ) -> None:
+        """
+        Add a cost entry to the tracker.
+        
+        Args:
+            provider: Provider name (e.g., 'openai', 'anthropic')
+            model: Model name
+            prompt_tokens: Number of prompt tokens
+            completion_tokens: Number of completion tokens
+            total_tokens: Total tokens used
+            cost_usd: Cost in USD
+            request_id: Unique request identifier
+            cached_tokens: Number of cached tokens
+            cache_creation_tokens: Tokens written to cache
+            cache_read_tokens: Tokens read from cache
+            group: Optional group tag for categorization
+        """
+        entry = CostEntry(
+            timestamp=datetime.now(),
+            provider=provider,
+            model=model,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=total_tokens,
+            cost_usd=cost_usd,
+            request_id=request_id,
+            cached_tokens=cached_tokens,
+            cache_creation_tokens=cache_creation_tokens,
+            cache_read_tokens=cache_read_tokens,
+            group=group,
+        )
+        self._entries.append(entry)
+        self._total_cost += cost_usd
+        
+        # Check budget alerts
+        if self._alert_threshold and self._total_cost >= self._alert_threshold:
+            self._trigger_alert(self._total_cost, self._alert_threshold)
+    
+    def get_total_cost(self) -> float:
+        """Get total cost across all tracked calls."""
+        return self._total_cost
+    
+    def get_total_tokens(self) -> int:
+        """Get total tokens across all tracked calls."""
+        return sum(entry.total_tokens for entry in self._entries)
+    
+    def get_call_count(self) -> int:
+        """Get total number of tracked calls."""
+        return len(self._entries)
+    
+    def get_entries(
+        self,
+        provider: Optional[str] = None,
+        model: Optional[str] = None,
+        group: Optional[str] = None,
+    ) -> List[CostEntry]:
+        """
+        Get filtered cost entries.
+        
+        Args:
+            provider: Filter by provider name
+            model: Filter by model name
+            group: Filter by group tag
+            
+        Returns:
+            List of matching cost entries
+        """
+        entries = self._entries
+        
+        if provider:
+            entries = [e for e in entries if e.provider == provider]
+        if model:
+            entries = [e for e in entries if e.model == model]
+        if group:
+            entries = [e for e in entries if e.group == group]
+        
+        return entries
+    
+    def get_cost_by_provider(self) -> Dict[str, float]:
+        """Get total cost grouped by provider."""
+        costs: Dict[str, float] = defaultdict(float)
+        for entry in self._entries:
+            costs[entry.provider] += entry.cost_usd
+        return dict(costs)
+    
+    def get_cost_by_model(self) -> Dict[str, float]:
+        """Get total cost grouped by model."""
+        costs: Dict[str, float] = defaultdict(float)
+        for entry in self._entries:
+            costs[entry.model] += entry.cost_usd
+        return dict(costs)
+    
+    def get_cost_by_group(self) -> Dict[str, float]:
+        """Get total cost grouped by custom group tag."""
+        costs: Dict[str, float] = defaultdict(float)
+        for entry in self._entries:
+            if entry.group:
+                costs[entry.group] += entry.cost_usd
+        return dict(costs)
+    
+    def get_tokens_by_provider(self) -> Dict[str, int]:
+        """Get total tokens grouped by provider."""
+        tokens: Dict[str, int] = defaultdict(int)
+        for entry in self._entries:
+            tokens[entry.provider] += entry.total_tokens
+        return dict(tokens)
+    
+    def get_cache_stats(self) -> Dict[str, any]:
+        """Get cache usage statistics."""
+        total_cache_reads = sum(e.cache_read_tokens for e in self._entries)
+        total_cache_creates = sum(e.cache_creation_tokens for e in self._entries)
+        total_prompt_tokens = sum(e.prompt_tokens for e in self._entries)
+        
+        cache_hit_rate = 0.0
+        if total_prompt_tokens > 0:
+            cache_hit_rate = (total_cache_reads / total_prompt_tokens) * 100
+        
+        return {
+            "total_cache_read_tokens": total_cache_reads,
+            "total_cache_creation_tokens": total_cache_creates,
+            "cache_hit_rate_percent": round(cache_hit_rate, 2),
+        }
+    
+    def set_budget(self, limit: float, alert_threshold: Optional[float] = None) -> None:
+        """
+        Set budget limit and optional alert threshold.
+        
+        Args:
+            limit: Maximum budget in USD
+            alert_threshold: Alert when cost reaches this threshold (default: 80% of limit)
+        """
+        self._budget_limit = limit
+        self._alert_threshold = alert_threshold or (limit * 0.8)
+    
+    def get_budget_status(self) -> Dict[str, any]:
+        """
+        Get current budget status.
+        
+        Returns:
+            Dictionary with budget information
+        """
+        if self._budget_limit is None:
+            return {
+                "budget_set": False,
+                "total_cost": self._total_cost,
+            }
+        
+        remaining = self._budget_limit - self._total_cost
+        percent_used = (self._total_cost / self._budget_limit) * 100
+        
+        return {
+            "budget_set": True,
+            "budget_limit": self._budget_limit,
+            "total_cost": self._total_cost,
+            "remaining": max(0, remaining),
+            "percent_used": round(percent_used, 2),
+            "over_budget": self._total_cost > self._budget_limit,
+            "alert_threshold": self._alert_threshold,
+        }
+    
+    def is_over_budget(self) -> bool:
+        """Check if current spending exceeds budget limit."""
+        if self._budget_limit is None:
+            return False
+        return self._total_cost > self._budget_limit
+    
+    def reset(self) -> None:
+        """Reset all tracked data."""
+        self._entries.clear()
+        self._total_cost = 0.0
+    
+    def _trigger_alert(self, current_cost: float, threshold: float) -> None:
+        """
+        Trigger budget alert (can be overridden for custom behavior).
+        
+        Args:
+            current_cost: Current total cost
+            threshold: Alert threshold that was exceeded
+        """
+        # Default implementation: print warning
+        # Override this method for custom alert behavior (email, webhook, etc.)
+        print(f"⚠️  Budget Alert: Current cost ${current_cost:.4f} exceeds threshold ${threshold:.4f}")
+    
+    def get_summary(self) -> Dict[str, any]:
+        """
+        Get comprehensive summary of tracked costs.
+        
+        Returns:
+            Dictionary with summary statistics
+        """
+        return {
+            "total_cost": self._total_cost,
+            "total_tokens": self.get_total_tokens(),
+            "total_calls": self.get_call_count(),
+            "cost_by_provider": self.get_cost_by_provider(),
+            "cost_by_model": self.get_cost_by_model(),
+            "tokens_by_provider": self.get_tokens_by_provider(),
+            "cache_stats": self.get_cache_stats(),
+            "budget_status": self.get_budget_status(),
+        }

stratifyai/embeddings.py

@@ -0,0 +1,245 @@
+"""Embedding generation for RAG and semantic search.
+
+This module provides abstraction for generating embeddings from text using
+various provider APIs (OpenAI, Cohere, etc.).
+"""
+
+import asyncio
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import List, Optional
+import os
+from openai import AsyncOpenAI
+
+from .exceptions import ProviderAPIError, AuthenticationError
+
+
+@dataclass
+class EmbeddingResult:
+    """Result of an embedding generation request.
+    
+    Attributes:
+        embeddings: List of embedding vectors (each is List[float])
+        model: Name of the embedding model used
+        total_tokens: Total tokens processed
+        cost: Cost of the embedding request in USD
+    """
+    embeddings: List[List[float]]
+    model: str
+    total_tokens: int
+    cost: float
+
+
+class EmbeddingProvider(ABC):
+    """Abstract base class for embedding providers.
+    
+    All embedding provider implementations must inherit from this class
+    and implement the generate_embeddings method.
+    """
+    
+    @abstractmethod
+    async def generate_embeddings(
+        self,
+        texts: List[str],
+        model: Optional[str] = None
+    ) -> EmbeddingResult:
+        """Generate embeddings for a list of texts.
+        
+        Args:
+            texts: List of text strings to embed
+            model: Optional model name override
+            
+        Returns:
+            EmbeddingResult with embeddings and metadata
+            
+        Raises:
+            ProviderAPIError: If the API request fails
+            AuthenticationError: If authentication fails
+        """
+        pass
+    
+    @abstractmethod
+    def get_embedding_dimension(self, model: str) -> int:
+        """Get the dimensionality of embeddings for a given model.
+        
+        Args:
+            model: Model name
+            
+        Returns:
+            Embedding dimension (e.g., 1536 for text-embedding-3-small)
+        """
+        pass
+    
+    def generate_embeddings_sync(
+        self,
+        texts: List[str],
+        model: Optional[str] = None
+    ) -> EmbeddingResult:
+        """Synchronous wrapper for generate_embeddings."""
+        return asyncio.run(self.generate_embeddings(texts, model))
+
+
+class OpenAIEmbeddingProvider(EmbeddingProvider):
+    """OpenAI embedding provider implementation.
+    
+    Supports:
+    - text-embedding-3-small (1536 dimensions)
+    - text-embedding-3-large (3072 dimensions)
+    - text-embedding-ada-002 (1536 dimensions, legacy)
+    """
+    
+    # Embedding costs per 1M tokens (as of Feb 2026)
+    EMBEDDING_COSTS = {
+        "text-embedding-3-small": 0.020 / 1_000_000,  # $0.020 per 1M tokens
+        "text-embedding-3-large": 0.130 / 1_000_000,  # $0.130 per 1M tokens
+        "text-embedding-ada-002": 0.100 / 1_000_000,  # $0.100 per 1M tokens (legacy)
+    }
+    
+    # Embedding dimensions by model
+    EMBEDDING_DIMENSIONS = {
+        "text-embedding-3-small": 1536,
+        "text-embedding-3-large": 3072,
+        "text-embedding-ada-002": 1536,
+    }
+    
+    # Default model
+    DEFAULT_MODEL = "text-embedding-3-small"
+    
+    def __init__(self, api_key: Optional[str] = None):
+        """Initialize OpenAI embedding provider.
+        
+        Args:
+            api_key: OpenAI API key. If None, reads from OPENAI_API_KEY env var.
+            
+        Raises:
+            AuthenticationError: If no API key is provided or found
+        """
+        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
+        if not self.api_key:
+            raise AuthenticationError(
+                "OpenAI API key not provided. Set OPENAI_API_KEY environment variable "
+                "or pass api_key parameter."
+            )
+        
+        self.client = AsyncOpenAI(api_key=self.api_key)
+    
+    async def generate_embeddings(
+        self,
+        texts: List[str],
+        model: Optional[str] = None
+    ) -> EmbeddingResult:
+        """Generate embeddings for a list of texts using OpenAI.
+        
+        Args:
+            texts: List of text strings to embed
+            model: Model name (default: text-embedding-3-small)
+            
+        Returns:
+            EmbeddingResult with embeddings and metadata
+            
+        Raises:
+            ProviderAPIError: If the API request fails
+            AuthenticationError: If authentication fails
+        """
+        model = model or self.DEFAULT_MODEL
+        
+        if model not in self.EMBEDDING_COSTS:
+            raise ValueError(
+                f"Unknown OpenAI embedding model: {model}. "
+                f"Supported models: {list(self.EMBEDDING_COSTS.keys())}"
+            )
+        
+        if not texts:
+            return EmbeddingResult(
+                embeddings=[],
+                model=model,
+                total_tokens=0,
+                cost=0.0
+            )
+        
+        try:
+            # Call OpenAI API
+            response = await self.client.embeddings.create(
+                input=texts,
+                model=model
+            )
+            
+            # Extract embeddings
+            embeddings = [data.embedding for data in response.data]
+            
+            # Calculate cost
+            total_tokens = response.usage.total_tokens
+            cost = total_tokens * self.EMBEDDING_COSTS[model]
+            
+            return EmbeddingResult(
+                embeddings=embeddings,
+                model=model,
+                total_tokens=total_tokens,
+                cost=cost
+            )
+            
+        except Exception as e:
+            error_msg = str(e)
+            if "authentication" in error_msg.lower() or "api key" in error_msg.lower():
+                raise AuthenticationError(f"OpenAI authentication failed: {error_msg}")
+            else:
+                raise ProviderAPIError(f"OpenAI embedding request failed: {error_msg}")
+    
+    async def generate_embedding(self, text: str, model: Optional[str] = None) -> List[float]:
+        """Generate embedding for a single text string.
+        
+        Convenience method for single text embedding.
+        
+        Args:
+            text: Text string to embed
+            model: Model name (default: text-embedding-3-small)
+            
+        Returns:
+            Embedding vector as List[float]
+        """
+        result = await self.generate_embeddings([text], model=model)
+        return result.embeddings[0]
+    
+    def get_embedding_dimension(self, model: str) -> int:
+        """Get the dimensionality of embeddings for a given model.
+        
+        Args:
+            model: Model name
+            
+        Returns:
+            Embedding dimension
+            
+        Raises:
+            ValueError: If model is unknown
+        """
+        if model not in self.EMBEDDING_DIMENSIONS:
+            raise ValueError(
+                f"Unknown OpenAI embedding model: {model}. "
+                f"Supported models: {list(self.EMBEDDING_DIMENSIONS.keys())}"
+            )
+        return self.EMBEDDING_DIMENSIONS[model]
+
+
+def create_embedding_provider(
+    provider: str = "openai",
+    api_key: Optional[str] = None
+) -> EmbeddingProvider:
+    """Factory function to create embedding providers.
+    
+    Args:
+        provider: Provider name (currently only "openai" supported)
+        api_key: API key for the provider
+        
+    Returns:
+        EmbeddingProvider instance
+        
+    Raises:
+        ValueError: If provider is unknown
+    """
+    if provider.lower() == "openai":
+        return OpenAIEmbeddingProvider(api_key=api_key)
+    else:
+        raise ValueError(
+            f"Unknown embedding provider: {provider}. "
+            f"Currently supported: openai"
+        )

stratifyai/exceptions.py

@@ -0,0 +1,91 @@
+"""Custom exceptions for LLM abstraction layer."""
+
+
+class LLMAbstractionError(Exception):
+    """Base exception for all LLM abstraction errors."""
+    pass
+
+
+class ProviderError(LLMAbstractionError):
+    """Base exception for provider-specific errors."""
+    pass
+
+
+class InvalidProviderError(ProviderError):
+    """Raised when an invalid provider is specified."""
+    pass
+
+
+class ProviderAPIError(ProviderError):
+    """Raised when a provider API call fails."""
+    
+    def __init__(self, message: str, provider: str, status_code: int = None):
+        self.provider = provider
+        self.status_code = status_code
+        super().__init__(f"[{provider}] {message}")
+
+
+class AuthenticationError(ProviderError):
+    """Raised when API key authentication fails."""
+    
+    def __init__(self, provider: str):
+        self.provider = provider
+        super().__init__(f"Authentication failed for {provider}. Check API key.")
+
+
+class InsufficientBalanceError(ProviderError):
+    """Raised when provider account has insufficient balance."""
+    
+    def __init__(self, provider: str):
+        self.provider = provider
+        super().__init__(f"Insufficient balance in {provider} account. Please add credits.")
+
+
+class RateLimitError(ProviderError):
+    """Raised when rate limit is exceeded."""
+    
+    def __init__(self, provider: str, retry_after: int = None):
+        self.provider = provider
+        self.retry_after = retry_after
+        message = f"Rate limit exceeded for {provider}"
+        if retry_after:
+            message += f". Retry after {retry_after} seconds"
+        super().__init__(message)
+
+
+class InvalidModelError(ProviderError):
+    """Raised when an invalid model is specified for a provider."""
+    
+    def __init__(self, model: str, provider: str):
+        self.model = model
+        self.provider = provider
+        super().__init__(f"Model '{model}' not supported by {provider}")
+
+
+class BudgetExceededError(LLMAbstractionError):
+    """Raised when budget limit is exceeded."""
+    
+    def __init__(self, current_cost: float, budget_limit: float):
+        self.current_cost = current_cost
+        self.budget_limit = budget_limit
+        super().__init__(
+            f"Budget limit ${budget_limit:.2f} exceeded. "
+            f"Current spend: ${current_cost:.2f}"
+        )
+
+
+class MaxRetriesExceededError(LLMAbstractionError):
+    """Raised when maximum retry attempts are exceeded."""
+    
+    def __init__(self, attempts: int, last_error: Exception):
+        self.attempts = attempts
+        self.last_error = last_error
+        super().__init__(
+            f"Maximum retry attempts ({attempts}) exceeded. "
+            f"Last error: {str(last_error)}"
+        )
+
+
+class ValidationError(LLMAbstractionError):
+    """Raised when input validation fails."""
+    pass

stratifyai/models.py

@@ -0,0 +1,59 @@
+"""Data models for unified LLM abstraction layer."""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import List, Literal, Optional
+
+
+@dataclass
+class Message:
+    """Standard message format for all providers (OpenAI-compatible)."""
+    role: Literal["system", "user", "assistant"]
+    content: str
+    name: Optional[str] = None  # For multi-agent scenarios
+    cache_control: Optional[dict] = None  # For providers that support prompt caching (Anthropic, OpenAI)
+
+
+@dataclass
+class Usage:
+    """Token usage and cost information."""
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    cached_tokens: int = 0  # Tokens retrieved from cache
+    cache_creation_tokens: int = 0  # Tokens written to cache (Anthropic)
+    cache_read_tokens: int = 0  # Tokens read from cache (Anthropic)
+    reasoning_tokens: int = 0  # For reasoning models like o1/o3
+    cost_usd: float = 0.0
+    cost_breakdown: Optional[dict] = None  # Detailed cost breakdown by token type
+
+
+@dataclass
+class ChatRequest:
+    """Unified request structure for chat completions."""
+    model: str
+    messages: List[Message]
+    temperature: float = 0.7
+    max_tokens: Optional[int] = None
+    stream: bool = False
+    top_p: float = 1.0
+    frequency_penalty: float = 0.0
+    presence_penalty: float = 0.0
+    stop: Optional[List[str]] = None
+    # Provider-specific extensions
+    reasoning_effort: Optional[str] = None  # OpenAI o1/o3
+    extra_params: dict = field(default_factory=dict)
+
+
+@dataclass
+class ChatResponse:
+    """Standard response from any provider."""
+    id: str
+    model: str
+    content: str
+    finish_reason: str
+    usage: Usage
+    provider: str
+    created_at: datetime
+    raw_response: dict  # Original provider response for debugging
+    latency_ms: Optional[float] = None  # Response latency in milliseconds

stratifyai/providers/__init__.py

@@ -0,0 +1,5 @@
+"""LLM provider implementations."""
+
+from .base import BaseProvider
+
+__all__ = ["BaseProvider"]