stratifyai 0.1.0__py3-none-any.whl

stratifyai/caching.py

@@ -0,0 +1,279 @@
+"""Caching utilities for LLM responses."""
+
+import hashlib
+import json
+import threading
+import time
+from dataclasses import dataclass, field
+from functools import wraps
+from typing import Any, Callable, Dict, Optional, Tuple
+
+from .models import ChatResponse
+
+
+@dataclass
+class CacheEntry:
+    """Entry in the response cache."""
+    response: ChatResponse
+    timestamp: float
+    hits: int = 0
+    cost_saved: float = 0.0  # Total cost saved from this entry
+
+
+class ResponseCache:
+    """Thread-safe in-memory cache for LLM responses."""
+    
+    def __init__(self, ttl: int = 3600, max_size: int = 1000):
+        """
+        Initialize response cache.
+        
+        Args:
+            ttl: Time-to-live for cache entries in seconds
+            max_size: Maximum number of entries to store
+        """
+        self.ttl = ttl
+        self.max_size = max_size
+        self._cache: Dict[str, CacheEntry] = {}
+        self._lock = threading.Lock()
+        self._total_misses: int = 0
+        self._total_cost_saved: float = 0.0
+    
+    def get(self, key: str) -> Optional[ChatResponse]:
+        """
+        Get response from cache.
+        
+        Args:
+            key: Cache key
+            
+        Returns:
+            Cached response if found and not expired, None otherwise
+        """
+        with self._lock:
+            if key not in self._cache:
+                self._total_misses += 1
+                return None
+            
+            entry = self._cache[key]
+            
+            # Check if expired
+            if time.time() - entry.timestamp > self.ttl:
+                del self._cache[key]
+                self._total_misses += 1
+                return None
+            
+            # Update hit count and cost saved
+            entry.hits += 1
+            if hasattr(entry.response, 'usage') and hasattr(entry.response.usage, 'cost_usd'):
+                cost = entry.response.usage.cost_usd
+                entry.cost_saved += cost
+                self._total_cost_saved += cost
+            
+            return entry.response
+    
+    def set(self, key: str, response: ChatResponse) -> None:
+        """
+        Store response in cache.
+        
+        Args:
+            key: Cache key
+            response: Response to cache
+        """
+        with self._lock:
+            # Evict oldest entry if cache is full
+            if len(self._cache) >= self.max_size:
+                oldest_key = min(
+                    self._cache.keys(),
+                    key=lambda k: self._cache[k].timestamp
+                )
+                del self._cache[oldest_key]
+            
+            self._cache[key] = CacheEntry(
+                response=response,
+                timestamp=time.time(),
+                hits=0
+            )
+    
+    def clear(self) -> None:
+        """Clear all cache entries."""
+        with self._lock:
+            self._cache.clear()
+            self._total_misses = 0
+            self._total_cost_saved = 0.0
+    
+    def get_stats(self) -> Dict[str, Any]:
+        """
+        Get cache statistics.
+        
+        Returns:
+            Dictionary with cache stats including hits, misses, and cost savings
+        """
+        with self._lock:
+            total_hits = sum(entry.hits for entry in self._cache.values())
+            total_requests = total_hits + self._total_misses
+            hit_rate = (total_hits / total_requests * 100) if total_requests > 0 else 0.0
+            
+            return {
+                "size": len(self._cache),
+                "max_size": self.max_size,
+                "total_hits": total_hits,
+                "total_misses": self._total_misses,
+                "total_requests": total_requests,
+                "hit_rate": hit_rate,
+                "total_cost_saved": self._total_cost_saved,
+                "ttl": self.ttl,
+            }
+    
+    def get_entries(self) -> list[Dict[str, Any]]:
+        """
+        Get detailed information about cache entries.
+        
+        Returns:
+            List of cache entry details
+        """
+        with self._lock:
+            entries = []
+            for key, entry in self._cache.items():
+                age = time.time() - entry.timestamp
+                entries.append({
+                    "key": key[:16] + "...",  # Truncate hash for display
+                    "model": entry.response.model if hasattr(entry.response, 'model') else "unknown",
+                    "provider": entry.response.provider if hasattr(entry.response, 'provider') else "unknown",
+                    "hits": entry.hits,
+                    "cost_saved": entry.cost_saved,
+                    "age_seconds": int(age),
+                    "expires_in": int(self.ttl - age),
+                })
+            
+            # Sort by hits (most popular first)
+            entries.sort(key=lambda x: x["hits"], reverse=True)
+            return entries
+
+
+# Global cache instance
+_global_cache = ResponseCache()
+
+
+def generate_cache_key(
+    model: str,
+    messages: list,
+    temperature: float,
+    max_tokens: Optional[int] = None,
+    **kwargs
+) -> str:
+    """
+    Generate a unique cache key from request parameters.
+    
+    Args:
+        model: Model name
+        messages: List of messages
+        temperature: Temperature parameter
+        max_tokens: Maximum tokens
+        **kwargs: Additional parameters
+        
+    Returns:
+        SHA256 hash of the request parameters
+    """
+    # Convert messages to hashable format
+    messages_str = json.dumps(
+        [{"role": m.role, "content": m.content} if hasattr(m, "role") else m 
+         for m in messages],
+        sort_keys=True
+    )
+    
+    # Include relevant parameters
+    cache_data = {
+        "model": model,
+        "messages": messages_str,
+        "temperature": temperature,
+        "max_tokens": max_tokens,
+    }
+    
+    # Add any additional kwargs that affect the response
+    for key in sorted(kwargs.keys()):
+        if key not in ["stream", "extra_params"]:  # Skip non-deterministic params
+            cache_data[key] = kwargs[key]
+    
+    # Generate hash
+    cache_str = json.dumps(cache_data, sort_keys=True)
+    return hashlib.sha256(cache_str.encode()).hexdigest()
+
+
+def cache_response(
+    ttl: int = 3600,
+    cache_instance: Optional[ResponseCache] = None
+):
+    """
+    Decorator to cache async LLM responses.
+    
+    Args:
+        ttl: Time-to-live for cache entries in seconds
+        cache_instance: Optional cache instance (uses global cache if None)
+        
+    Returns:
+        Decorated async function
+        
+    Example:
+        @cache_response(ttl=3600)
+        async def chat(self, request: ChatRequest) -> ChatResponse:
+            return await self.provider.chat_completion(request)
+    """
+    cache = cache_instance or _global_cache
+    cache.ttl = ttl
+    
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs) -> ChatResponse:
+            # Extract request parameters
+            # Handle both ChatRequest object and individual parameters
+            if args and hasattr(args[0], "model"):
+                request = args[0]
+                cache_key = generate_cache_key(
+                    model=request.model,
+                    messages=request.messages,
+                    temperature=request.temperature,
+                    max_tokens=request.max_tokens,
+                )
+            else:
+                cache_key = generate_cache_key(**kwargs)
+            
+            # Check cache
+            cached_response = cache.get(cache_key)
+            if cached_response is not None:
+                return cached_response
+            
+            # Execute async function
+            response = await func(*args, **kwargs)
+            
+            # Cache response (only if not streaming)
+            if not kwargs.get("stream", False):
+                cache.set(cache_key, response)
+            
+            return response
+        
+        return async_wrapper
+    return decorator
+
+
+def get_cache_stats() -> Dict[str, Any]:
+    """
+    Get statistics from the global cache.
+    
+    Returns:
+        Dictionary with cache statistics
+    """
+    return _global_cache.get_stats()
+
+
+def get_cache_entries() -> list[Dict[str, Any]]:
+    """
+    Get detailed information about cache entries.
+    
+    Returns:
+        List of cache entry details
+    """
+    return _global_cache.get_entries()
+
+
+def clear_cache() -> None:
+    """Clear the global cache."""
+    _global_cache.clear()

stratifyai/chat/__init__.py

@@ -0,0 +1,54 @@
+"""Chat package for StratifyAI provider-specific chat interfaces.
+
+This package provides convenient, provider-specific chat functions.
+Model must be specified for each request.
+
+Usage:
+    # Model is always required
+    from stratifyai.chat import openai, anthropic, google
+    response = await openai.chat("Hello!", model="gpt-4.1-mini")
+    
+    # Builder pattern (model required first)
+    from stratifyai.chat import anthropic
+    client = (
+        anthropic
+        .with_model("claude-opus-4-5")
+        .with_system("You are a helpful assistant")
+        .with_developer("Use markdown formatting")
+    )
+    response = await client.chat("Hello!")
+    
+    # With additional parameters
+    response = await anthropic.chat(
+        "Explain quantum computing",
+        model="claude-sonnet-4-5",
+        temperature=0.5,
+        max_tokens=500,
+    )
+"""
+
+from stratifyai.chat.builder import ChatBuilder
+from stratifyai.chat import (
+    stratifyai_openai as openai,
+    stratifyai_anthropic as anthropic,
+    stratifyai_google as google,
+    stratifyai_deepseek as deepseek,
+    stratifyai_groq as groq,
+    stratifyai_grok as grok,
+    stratifyai_openrouter as openrouter,
+    stratifyai_ollama as ollama,
+    stratifyai_bedrock as bedrock,
+)
+
+__all__ = [
+    "ChatBuilder",
+    "openai",
+    "anthropic",
+    "google",
+    "deepseek",
+    "groq",
+    "grok",
+    "openrouter",
+    "ollama",
+    "bedrock",
+]

stratifyai/chat/builder.py

@@ -0,0 +1,366 @@
+"""Builder pattern for configuring chat clients.
+
+Provides a fluent interface for configuring chat parameters before execution.
+
+Example:
+    from stratifyai.chat import anthropic
+    
+    # Model is always required
+    response = await anthropic.chat("Hello", model="claude-sonnet-4-5")
+    
+    # Builder pattern with chaining (model required)
+    client = (
+        anthropic
+        .with_model("claude-opus-4-5")
+        .with_system("I am a helpful assistant")
+        .with_developer("Use markdown formatting")
+        .with_temperature(0.5)
+        .with_max_tokens(1000)
+    )
+    response = await client.chat("Hello")
+"""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, AsyncIterator, Callable, Optional, Union
+
+if TYPE_CHECKING:
+    from stratifyai import LLMClient
+    from stratifyai.models import ChatResponse, Message
+
+
+@dataclass
+class ChatBuilder:
+    """Builder for configuring chat requests with fluent interface.
+    
+    Each with_* method returns a new ChatBuilder instance with the
+    updated configuration, allowing for method chaining.
+    
+    Model is required - either via with_model() or as a parameter to chat().
+    
+    Attributes:
+        provider: The provider name (e.g., "openai", "anthropic")
+        default_temperature: Default temperature setting
+        default_max_tokens: Default max tokens setting
+        _model: Configured model (required before chat)
+        _system: System prompt
+        _developer: Developer/instruction prompt
+        _temperature: Configured temperature (None = use default)
+        _max_tokens: Configured max tokens (None = use default)
+        _client_factory: Factory function to create LLMClient
+    """
+    provider: str
+    default_temperature: float = 0.7
+    default_max_tokens: Optional[int] = None
+    _model: Optional[str] = None
+    _system: Optional[str] = None
+    _developer: Optional[str] = None
+    _temperature: Optional[float] = None
+    _max_tokens: Optional[int] = None
+    _client_factory: Optional[Callable[[], "LLMClient"]] = None
+    _extra_kwargs: dict = field(default_factory=dict)
+    
+    def _clone(self, **updates) -> "ChatBuilder":
+        """Create a copy of this builder with updates applied."""
+        return ChatBuilder(
+            provider=self.provider,
+            default_temperature=self.default_temperature,
+            default_max_tokens=self.default_max_tokens,
+            _model=updates.get("_model", self._model),
+            _system=updates.get("_system", self._system),
+            _developer=updates.get("_developer", self._developer),
+            _temperature=updates.get("_temperature", self._temperature),
+            _max_tokens=updates.get("_max_tokens", self._max_tokens),
+            _client_factory=self._client_factory,
+            _extra_kwargs={**self._extra_kwargs, **updates.get("_extra_kwargs", {})},
+        )
+    
+    def with_model(self, model: str) -> "ChatBuilder":
+        """Set the model to use (required).
+        
+        Args:
+            model: Model name (e.g., "claude-opus-4-5", "gpt-4.1")
+            
+        Returns:
+            New ChatBuilder with model configured.
+            
+        Example:
+            client = anthropic.with_model("claude-opus-4-5")
+        """
+        return self._clone(_model=model)
+    
+    def with_system(self, prompt: str) -> "ChatBuilder":
+        """Set the system prompt.
+        
+        The system prompt sets the AI's behavior and context.
+        
+        Args:
+            prompt: System prompt text
+            
+        Returns:
+            New ChatBuilder with system prompt configured.
+            
+        Example:
+            client = anthropic.with_system("You are a helpful coding assistant")
+        """
+        return self._clone(_system=prompt)
+    
+    def with_developer(self, instructions: str) -> "ChatBuilder":
+        """Set developer/instruction prompt.
+        
+        Developer instructions provide formatting or behavioral guidance.
+        These are prepended to the system prompt if both are set.
+        
+        Args:
+            instructions: Developer instructions text
+            
+        Returns:
+            New ChatBuilder with developer instructions configured.
+            
+        Example:
+            client = anthropic.with_developer("Always use markdown formatting")
+        """
+        return self._clone(_developer=instructions)
+    
+    def with_temperature(self, temperature: float) -> "ChatBuilder":
+        """Set the sampling temperature.
+        
+        Args:
+            temperature: Temperature value (0.0 = deterministic, 2.0 = creative)
+            
+        Returns:
+            New ChatBuilder with temperature configured.
+            
+        Example:
+            client = anthropic.with_temperature(0.3)
+        """
+        return self._clone(_temperature=temperature)
+    
+    def with_max_tokens(self, max_tokens: int) -> "ChatBuilder":
+        """Set the maximum tokens to generate.
+        
+        Args:
+            max_tokens: Maximum number of tokens in the response
+            
+        Returns:
+            New ChatBuilder with max_tokens configured.
+            
+        Example:
+            client = anthropic.with_max_tokens(500)
+        """
+        return self._clone(_max_tokens=max_tokens)
+    
+    def with_options(self, **kwargs) -> "ChatBuilder":
+        """Set additional provider-specific options.
+        
+        Args:
+            **kwargs: Additional options passed to the API
+            
+        Returns:
+            New ChatBuilder with extra options configured.
+            
+        Example:
+            client = anthropic.with_options(top_p=0.9)
+        """
+        return self._clone(_extra_kwargs=kwargs)
+    
+    @property
+    def model(self) -> Optional[str]:
+        """Get the configured model (None if not set)."""
+        return self._model
+    
+    @property
+    def temperature(self) -> float:
+        """Get the effective temperature (configured or default)."""
+        return self._temperature if self._temperature is not None else self.default_temperature
+    
+    @property
+    def max_tokens(self) -> Optional[int]:
+        """Get the effective max_tokens (configured or default)."""
+        return self._max_tokens if self._max_tokens is not None else self.default_max_tokens
+    
+    def _get_client(self) -> "LLMClient":
+        """Get or create the LLM client."""
+        if self._client_factory is None:
+            from stratifyai import LLMClient
+            return LLMClient(provider=self.provider)
+        return self._client_factory()
+    
+    def _build_system_prompt(self) -> Optional[str]:
+        """Combine developer and system prompts."""
+        if self._developer and self._system:
+            return f"{self._developer}\n\n{self._system}"
+        return self._developer or self._system
+    
+    def _build_messages(self, prompt: Union[str, list]) -> list:
+        """Build the messages list from prompt and configured prompts."""
+        from stratifyai.models import Message
+        
+        if isinstance(prompt, str):
+            messages = []
+            system_prompt = self._build_system_prompt()
+            if system_prompt:
+                messages.append(Message(role="system", content=system_prompt))
+            messages.append(Message(role="user", content=prompt))
+            return messages
+        else:
+            # If prompt is already a list of messages, prepend system if not present
+            messages = list(prompt)
+            system_prompt = self._build_system_prompt()
+            if system_prompt and (not messages or messages[0].role != "system"):
+                from stratifyai.models import Message
+                messages.insert(0, Message(role="system", content=system_prompt))
+            return messages
+    
+    async def chat(
+        self,
+        prompt: Union[str, list],
+        *,
+        model: Optional[str] = None,
+        system: Optional[str] = None,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        stream: bool = False,
+        **kwargs,
+    ) -> Union["ChatResponse", AsyncIterator["ChatResponse"]]:
+        """Send an async chat completion request.
+        
+        Args:
+            prompt: User message string or list of Message objects.
+            model: Model to use (required if not set via with_model()).
+            system: Override the configured system prompt.
+            temperature: Override the configured temperature.
+            max_tokens: Override the configured max_tokens.
+            stream: Whether to stream the response.
+            **kwargs: Additional parameters passed to the API.
+            
+        Returns:
+            ChatResponse object, or AsyncIterator[ChatResponse] if streaming.
+            
+        Raises:
+            ValueError: If no model is specified.
+        """
+        client = self._get_client()
+        
+        # Use overrides if provided, otherwise use configured values
+        effective_model = model or self._model
+        if not effective_model:
+            raise ValueError(
+                f"Model is required. Either call with_model() first or pass model parameter.\n"
+                f"Example: {self.provider}.with_model('model-name').chat(...) or "
+                f"{self.provider}.chat(..., model='model-name')"
+            )
+        effective_temp = temperature if temperature is not None else self.temperature
+        effective_max = max_tokens if max_tokens is not None else self.max_tokens
+        
+        # Handle system prompt override
+        if system:
+            builder = self.with_system(system)
+            messages = builder._build_messages(prompt)
+        else:
+            messages = self._build_messages(prompt)
+        
+        # Merge extra kwargs
+        merged_kwargs = {**self._extra_kwargs, **kwargs}
+        
+        return await client.chat(
+            model=effective_model,
+            messages=messages,
+            temperature=effective_temp,
+            max_tokens=effective_max,
+            stream=stream,
+            **merged_kwargs,
+        )
+    
+    async def chat_stream(
+        self,
+        prompt: Union[str, list],
+        *,
+        model: Optional[str] = None,
+        system: Optional[str] = None,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        **kwargs,
+    ) -> AsyncIterator["ChatResponse"]:
+        """Send an async streaming chat completion request.
+        
+        Args:
+            prompt: User message string or list of Message objects.
+            model: Override the configured model.
+            system: Override the configured system prompt.
+            temperature: Override the configured temperature.
+            max_tokens: Override the configured max_tokens.
+            **kwargs: Additional parameters passed to the API.
+            
+        Yields:
+            ChatResponse chunks.
+        """
+        return await self.chat(
+            prompt,
+            model=model,
+            system=system,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            stream=True,
+            **kwargs,
+        )
+    
+    def chat_sync(
+        self,
+        prompt: Union[str, list],
+        *,
+        model: Optional[str] = None,
+        system: Optional[str] = None,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        **kwargs,
+    ) -> "ChatResponse":
+        """Synchronous wrapper for chat().
+        
+        Args:
+            prompt: User message string or list of Message objects.
+            model: Override the configured model.
+            system: Override the configured system prompt.
+            temperature: Override the configured temperature.
+            max_tokens: Override the configured max_tokens.
+            **kwargs: Additional parameters passed to the API.
+            
+        Returns:
+            ChatResponse object.
+        """
+        return asyncio.run(self.chat(
+            prompt,
+            model=model,
+            system=system,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            stream=False,
+            **kwargs,
+        ))
+
+
+def create_module_builder(
+    provider: str,
+    default_temperature: float = 0.7,
+    default_max_tokens: Optional[int] = None,
+    client_factory: Optional[Callable[[], "LLMClient"]] = None,
+) -> ChatBuilder:
+    """Create a ChatBuilder for a provider module.
+    
+    Args:
+        provider: Provider name
+        default_temperature: Default temperature
+        default_max_tokens: Default max tokens
+        client_factory: Optional factory to create shared client
+        
+    Returns:
+        Configured ChatBuilder instance
+    """
+    return ChatBuilder(
+        provider=provider,
+        default_temperature=default_temperature,
+        default_max_tokens=default_max_tokens,
+        _client_factory=client_factory,
+    )