stratifyai 0.1.0__py3-none-any.whl

stratifyai/providers/google.py

@@ -0,0 +1,39 @@
+"""Google Gemini provider implementation."""
+
+import os
+from typing import Optional
+
+from ..config import GOOGLE_MODELS, PROVIDER_BASE_URLS
+from ..exceptions import AuthenticationError
+from .openai_compatible import OpenAICompatibleProvider
+
+
+class GoogleProvider(OpenAICompatibleProvider):
+    """Google Gemini provider using OpenAI-compatible API."""
+    
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        config: dict = None
+    ):
+        """
+        Initialize Google Gemini provider.
+        
+        Args:
+            api_key: Google API key (defaults to GOOGLE_API_KEY env var)
+            config: Optional provider-specific configuration
+            
+        Raises:
+            AuthenticationError: If API key not provided
+        """
+        api_key = api_key or os.getenv("GOOGLE_API_KEY")
+        if not api_key:
+            raise AuthenticationError("google")
+        
+        base_url = PROVIDER_BASE_URLS["google"]
+        super().__init__(api_key, base_url, GOOGLE_MODELS, config)
+    
+    @property
+    def provider_name(self) -> str:
+        """Return provider name."""
+        return "google"

stratifyai/providers/grok.py

@@ -0,0 +1,39 @@
+"""Grok (X.AI) provider implementation."""
+
+import os
+from typing import Optional
+
+from ..config import GROK_MODELS, PROVIDER_BASE_URLS
+from ..exceptions import AuthenticationError
+from .openai_compatible import OpenAICompatibleProvider
+
+
+class GrokProvider(OpenAICompatibleProvider):
+    """Grok (X.AI) provider using OpenAI-compatible API."""
+    
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        config: dict = None
+    ):
+        """
+        Initialize Grok provider.
+        
+        Args:
+            api_key: Grok API key (defaults to GROK_API_KEY env var)
+            config: Optional provider-specific configuration
+            
+        Raises:
+            AuthenticationError: If API key not provided
+        """
+        api_key = api_key or os.getenv("GROK_API_KEY")
+        if not api_key:
+            raise AuthenticationError("grok")
+        
+        base_url = PROVIDER_BASE_URLS["grok"]
+        super().__init__(api_key, base_url, GROK_MODELS, config)
+    
+    @property
+    def provider_name(self) -> str:
+        """Return provider name."""
+        return "grok"

stratifyai/providers/groq.py

@@ -0,0 +1,39 @@
+"""Groq provider implementation."""
+
+import os
+from typing import Optional
+
+from ..config import GROQ_MODELS, PROVIDER_BASE_URLS
+from ..exceptions import AuthenticationError
+from .openai_compatible import OpenAICompatibleProvider
+
+
+class GroqProvider(OpenAICompatibleProvider):
+    """Groq provider using OpenAI-compatible API."""
+    
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        config: dict = None
+    ):
+        """
+        Initialize Groq provider.
+        
+        Args:
+            api_key: Groq API key (defaults to GROQ_API_KEY env var)
+            config: Optional provider-specific configuration
+            
+        Raises:
+            AuthenticationError: If API key not provided
+        """
+        api_key = api_key or os.getenv("GROQ_API_KEY")
+        if not api_key:
+            raise AuthenticationError("groq")
+        
+        base_url = PROVIDER_BASE_URLS["groq"]
+        super().__init__(api_key, base_url, GROQ_MODELS, config)
+    
+    @property
+    def provider_name(self) -> str:
+        """Return provider name."""
+        return "groq"

stratifyai/providers/ollama.py

@@ -0,0 +1,43 @@
+"""Ollama provider implementation for local models."""
+
+import os
+from typing import Optional
+
+from ..config import OLLAMA_MODELS, PROVIDER_BASE_URLS
+from .openai_compatible import OpenAICompatibleProvider
+
+
+class OllamaProvider(OpenAICompatibleProvider):
+    """Ollama provider for local models using OpenAI-compatible API."""
+    
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        config: dict = None
+    ):
+        """
+        Initialize Ollama provider.
+        
+        Args:
+            api_key: Optional API key (Ollama typically doesn't require one)
+            config: Optional provider-specific configuration (can include base_url)
+            
+        Note:
+            Ollama runs locally and typically doesn't require an API key.
+            Default base URL is http://localhost:11434/v1
+        """
+        # Ollama doesn't require an API key, use placeholder
+        from ..api_key_helper import APIKeyHelper
+        api_key = APIKeyHelper.get_api_key("ollama", api_key) or "ollama"
+        
+        # Allow custom base URL via config
+        base_url = PROVIDER_BASE_URLS["ollama"]
+        if config and "base_url" in config:
+            base_url = config["base_url"]
+        
+        super().__init__(api_key, base_url, OLLAMA_MODELS, config)
+    
+    @property
+    def provider_name(self) -> str:
+        """Return provider name."""
+        return "ollama"

stratifyai/providers/openai.py

@@ -0,0 +1,344 @@
+"""OpenAI provider implementation."""
+
+import os
+from datetime import datetime
+from typing import AsyncIterator, List, Optional
+
+from openai import AsyncOpenAI
+
+from ..config import OPENAI_MODELS, PROVIDER_CONSTRAINTS
+from ..exceptions import AuthenticationError, InvalidModelError, ProviderAPIError
+from ..models import ChatRequest, ChatResponse, Usage
+from .base import BaseProvider
+
+
+class OpenAIProvider(BaseProvider):
+    """OpenAI provider implementation with cost tracking."""
+    
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        config: dict = None
+    ):
+        """
+        Initialize OpenAI provider.
+        
+        Args:
+            api_key: OpenAI API key (defaults to OPENAI_API_KEY env var)
+            config: Optional provider-specific configuration
+            
+        Raises:
+            ValueError: If API key not provided (with helpful setup instructions)
+        """
+        from ..api_key_helper import get_api_key_or_error
+        api_key = get_api_key_or_error("openai", api_key)
+        super().__init__(api_key, config)
+        self._initialize_client()
+    
+    def _initialize_client(self) -> None:
+        """Initialize OpenAI async client."""
+        try:
+            self._client = AsyncOpenAI(api_key=self.api_key)
+        except Exception as e:
+            raise ProviderAPIError(
+                f"Failed to initialize OpenAI client: {str(e)}",
+                "openai"
+            )
+    
+    @property
+    def provider_name(self) -> str:
+        """Return provider name."""
+        return "openai"
+    
+    def get_supported_models(self) -> List[str]:
+        """Return list of supported OpenAI models."""
+        return list(OPENAI_MODELS.keys())
+    
+    def supports_caching(self, model: str) -> bool:
+        """Check if model supports prompt caching."""
+        model_info = OPENAI_MODELS.get(model, {})
+        return model_info.get("supports_caching", False)
+    
+    async def chat_completion(self, request: ChatRequest) -> ChatResponse:
+        """
+        Execute chat completion request.
+        
+        Args:
+            request: Unified chat request
+            
+        Returns:
+            Unified chat response with cost tracking
+            
+        Raises:
+            InvalidModelError: If model not supported
+            ProviderAPIError: If API call fails
+        """
+        if not self.validate_model(request.model):
+            raise InvalidModelError(request.model, self.provider_name)
+        
+        # Validate temperature constraints for OpenAI (0.0 to 2.0)
+        constraints = PROVIDER_CONSTRAINTS.get(self.provider_name, {})
+        self.validate_temperature(
+            request.temperature,
+            constraints.get("min_temperature", 0.0),
+            constraints.get("max_temperature", 2.0)
+        )
+        
+        # Build OpenAI-specific request parameters
+        messages = []
+        for msg in request.messages:
+            message_dict = {"role": msg.role, "content": msg.content}
+            # Add cache_control if present and model supports caching
+            if msg.cache_control and self.supports_caching(request.model):
+                message_dict["cache_control"] = msg.cache_control
+            messages.append(message_dict)
+        
+        openai_params = {
+            "model": request.model,
+            "messages": messages,
+        }
+        
+        # Check if model is a reasoning model (o-series)
+        model_info = OPENAI_MODELS.get(request.model, {})
+        is_reasoning_model = model_info.get("reasoning_model", False)
+        
+        # Also check if model name starts with o1, o3, gpt-5, or just 'o' followed by a digit
+        # This catches variants like o1-preview, o1-2024-12-17, o3-mini, gpt-5, etc.
+        if not is_reasoning_model and request.model:
+            model_lower = request.model.lower()
+            # Match: o1*, o3*, gpt-5*, "reasoning", or o followed by digit
+            is_reasoning_model = (
+                model_lower.startswith("o1") or 
+                model_lower.startswith("o3") or
+                model_lower.startswith("gpt-5") or
+                "reasoning" in model_lower or
+                (model_lower.startswith("o") and len(model_lower) > 1 and model_lower[1].isdigit())
+            )
+        
+        # Only add these parameters for non-reasoning models
+        # Reasoning models like o1, o1-mini, o3-mini don't support temperature/top_p/penalties
+        if not is_reasoning_model:
+            openai_params["temperature"] = request.temperature
+            openai_params["top_p"] = request.top_p
+            openai_params["frequency_penalty"] = request.frequency_penalty
+            openai_params["presence_penalty"] = request.presence_penalty
+        
+        # Add optional parameters
+        if request.max_tokens:
+            openai_params["max_tokens"] = request.max_tokens
+        if request.stop:
+            openai_params["stop"] = request.stop
+        
+        # Add reasoning_effort for o-series models
+        if request.reasoning_effort and "o" in request.model:
+            openai_params["reasoning_effort"] = request.reasoning_effort
+        
+        # Add any extra params
+        if request.extra_params:
+            openai_params.update(request.extra_params)
+        
+        try:
+            # Make API request
+            raw_response = await self._client.chat.completions.create(**openai_params)
+            # Normalize and return
+            return self._normalize_response(raw_response.model_dump())
+        except Exception as e:
+            raise ProviderAPIError(
+                f"Chat completion failed: {str(e)}",
+                self.provider_name
+            )
+    
+    async def chat_completion_stream(
+        self, request: ChatRequest
+    ) -> AsyncIterator[ChatResponse]:
+        """
+        Execute streaming chat completion request.
+        
+        Args:
+            request: Unified chat request
+            
+        Yields:
+            Unified chat response chunks
+            
+        Raises:
+            InvalidModelError: If model not supported
+            ProviderAPIError: If API call fails
+        """
+        if not self.validate_model(request.model):
+            raise InvalidModelError(request.model, self.provider_name)
+        
+        # Build request parameters
+        openai_params = {
+            "model": request.model,
+            "messages": [
+                {"role": msg.role, "content": msg.content}
+                for msg in request.messages
+            ],
+            "stream": True,
+        }
+        
+        # Check if model is a reasoning model
+        model_info = OPENAI_MODELS.get(request.model, {})
+        is_reasoning_model = model_info.get("reasoning_model", False)
+        
+        # Also check if model name starts with o1, o3, gpt-5, or just 'o' followed by a digit
+        if not is_reasoning_model and request.model:
+            model_lower = request.model.lower()
+            is_reasoning_model = (
+                model_lower.startswith("o1") or 
+                model_lower.startswith("o3") or
+                model_lower.startswith("gpt-5") or
+                "reasoning" in model_lower or
+                (model_lower.startswith("o") and len(model_lower) > 1 and model_lower[1].isdigit())
+            )
+        
+        # Only add temperature for non-reasoning models
+        if not is_reasoning_model:
+            openai_params["temperature"] = request.temperature
+        
+        if request.max_tokens:
+            openai_params["max_tokens"] = request.max_tokens
+        
+        try:
+            stream = await self._client.chat.completions.create(**openai_params)
+            
+            async for chunk in stream:
+                chunk_dict = chunk.model_dump()
+                if chunk.choices and chunk.choices[0].delta.content:
+                    yield self._normalize_stream_chunk(chunk_dict)
+        except Exception as e:
+            raise ProviderAPIError(
+                f"Streaming chat completion failed: {str(e)}",
+                self.provider_name
+            )
+    
+    def _normalize_response(self, raw_response: dict) -> ChatResponse:
+        """
+        Convert OpenAI response to unified format.
+        
+        Args:
+            raw_response: Raw OpenAI API response
+            
+        Returns:
+            Normalized ChatResponse with cost
+        """
+        choice = raw_response["choices"][0]
+        usage_dict = raw_response.get("usage", {})
+        
+        # Extract token usage
+        prompt_details = usage_dict.get("prompt_tokens_details", {})
+        usage = Usage(
+            prompt_tokens=usage_dict.get("prompt_tokens", 0),
+            completion_tokens=usage_dict.get("completion_tokens", 0),
+            total_tokens=usage_dict.get("total_tokens", 0),
+            cached_tokens=prompt_details.get("cached_tokens", 0),
+            cache_creation_tokens=prompt_details.get("cache_creation_input_tokens", 0),
+            cache_read_tokens=prompt_details.get("cached_tokens", 0),
+            reasoning_tokens=usage_dict.get("completion_tokens_details", {}).get(
+                "reasoning_tokens", 0
+            ),
+        )
+        
+        # Calculate cost including cache costs
+        base_cost = self._calculate_cost(usage, raw_response["model"])
+        cache_cost = self._calculate_cache_cost(
+            usage.cache_creation_tokens,
+            usage.cache_read_tokens,
+            raw_response["model"]
+        )
+        usage.cost_usd = base_cost + cache_cost
+        
+        # Add cost breakdown
+        if usage.cache_creation_tokens > 0 or usage.cache_read_tokens > 0:
+            usage.cost_breakdown = {
+                "base_cost": base_cost,
+                "cache_cost": cache_cost,
+                "total_cost": usage.cost_usd,
+            }
+        
+        return ChatResponse(
+            id=raw_response["id"],
+            model=raw_response["model"],
+            content=choice["message"]["content"] or "",
+            finish_reason=choice["finish_reason"],
+            usage=usage,
+            provider=self.provider_name,
+            created_at=datetime.fromtimestamp(raw_response["created"]),
+            raw_response=raw_response,
+        )
+    
+    def _normalize_stream_chunk(self, chunk_dict: dict) -> ChatResponse:
+        """Normalize streaming chunk to ChatResponse format."""
+        choice = chunk_dict["choices"][0]
+        content = choice["delta"].get("content", "")
+        
+        return ChatResponse(
+            id=chunk_dict["id"],
+            model=chunk_dict["model"],
+            content=content,
+            finish_reason=choice.get("finish_reason", ""),
+            usage=Usage(
+                prompt_tokens=0,
+                completion_tokens=0,
+                total_tokens=0
+            ),
+            provider=self.provider_name,
+            created_at=datetime.fromtimestamp(chunk_dict["created"]),
+            raw_response=chunk_dict,
+        )
+    
+    def _calculate_cost(self, usage: Usage, model: str) -> float:
+        """
+        Calculate cost in USD based on token usage (excluding cache costs).
+        
+        Args:
+            usage: Token usage information
+            model: Model name used
+            
+        Returns:
+            Cost in USD
+        """
+        model_info = OPENAI_MODELS.get(model, {})
+        cost_input = model_info.get("cost_input", 0.0)
+        cost_output = model_info.get("cost_output", 0.0)
+        
+        # Calculate non-cached prompt tokens
+        non_cached_prompt_tokens = usage.prompt_tokens - usage.cache_read_tokens
+        
+        # Costs are per 1M tokens
+        input_cost = (non_cached_prompt_tokens / 1_000_000) * cost_input
+        output_cost = (usage.completion_tokens / 1_000_000) * cost_output
+        
+        return input_cost + output_cost
+    
+    def _calculate_cache_cost(
+        self,
+        cache_creation_tokens: int,
+        cache_read_tokens: int,
+        model: str
+    ) -> float:
+        """
+        Calculate cost for cached tokens.
+        
+        Args:
+            cache_creation_tokens: Number of tokens written to cache
+            cache_read_tokens: Number of tokens read from cache
+            model: Model name used
+            
+        Returns:
+            Cost in USD for cache operations
+        """
+        model_info = OPENAI_MODELS.get(model, {})
+        
+        # Check if model supports caching
+        if not model_info.get("supports_caching", False):
+            return 0.0
+        
+        cost_cache_write = model_info.get("cost_cache_write", 0.0)
+        cost_cache_read = model_info.get("cost_cache_read", 0.0)
+        
+        # Costs are per 1M tokens
+        write_cost = (cache_creation_tokens / 1_000_000) * cost_cache_write
+        read_cost = (cache_read_tokens / 1_000_000) * cost_cache_read
+        
+        return write_cost + read_cost