stratifyai 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

stratifyai/api_key_helper.py

@@ -15,7 +15,7 @@ class APIKeyHelper:
         "google": "GOOGLE_API_KEY",
         "deepseek": "DEEPSEEK_API_KEY",
         "groq": "GROQ_API_KEY",
-        "grok": "GROK_API_KEY",
+        "grok": "XAI_API_KEY",  # X.AI official environment variable
         "openrouter": "OPENROUTER_API_KEY",
         "ollama": "OLLAMA_API_KEY",
         "bedrock": "AWS_BEARER_TOKEN_BEDROCK",  # Bedrock bearer token (or AWS_ACCESS_KEY_ID)

stratifyai/config.py

@@ -29,7 +29,7 @@ OPENAI_MODELS: Dict[str, Dict[str, Any]] = {
         "context": 128000,
         "cost_input": 10.0,
         "cost_output": 30.0,
-        "supports_vision": True,
+        "supports_vision": False,
         "supports_tools": True,
     },
     "gpt-4": {
@@ -310,16 +310,6 @@ ANTHROPIC_MODELS: Dict[str, Dict[str, Any]] = {
         "supports_tools": True,
         "supports_caching": True,
     },
-    "claude-3-5-haiku-20241022": {
-        "context": 200000,
-        "cost_input": 1.0,
-        "cost_output": 5.0,
-        "cost_cache_write": 1.25,
-        "cost_cache_read": 0.10,
-        "supports_vision": False,
-        "supports_tools": True,
-        "supports_caching": True,
-    },
 }
 
 # Google Gemini Model Catalog (OpenAI-compatible)
@@ -449,6 +439,105 @@ GROQ_MODELS: Dict[str, Dict[str, Any]] = {
 
 # Grok (X.AI) Model Catalog (OpenAI-compatible)
 GROK_MODELS: Dict[str, Dict[str, Any]] = {
+    # Grok 4.1 Fast - Best for agentic tool calling
+    "grok-4-1-fast-reasoning": {
+        "context": 2000000,  # 2M tokens
+        "cost_input": 0.20,
+        "cost_output": 0.50,
+        "supports_vision": True,
+        "supports_tools": True,
+        "reasoning_model": True,
+    },
+    "grok-4-1-fast-non-reasoning": {
+        "context": 2000000,
+        "cost_input": 0.20,
+        "cost_output": 0.50,
+        "supports_vision": True,
+        "supports_tools": True,
+    },
+    # Grok 4 Fast - High performance
+    "grok-4-fast-reasoning": {
+        "context": 2000000,
+        "cost_input": 0.50,
+        "cost_output": 1.50,
+        "supports_vision": True,
+        "supports_tools": True,
+        "reasoning_model": True,
+    },
+    "grok-4-fast-non-reasoning": {
+        "context": 2000000,
+        "cost_input": 0.50,
+        "cost_output": 1.50,
+        "supports_vision": True,
+        "supports_tools": True,
+    },
+    # Grok 4 - Flagship reasoning model
+    "grok-4": {
+        "context": 128000,
+        "cost_input": 5.0,
+        "cost_output": 15.0,
+        "supports_vision": True,
+        "supports_tools": True,
+        "reasoning_model": True,
+        "fixed_temperature": 1.0,
+    },
+    # Grok Code Fast - Specialized for coding
+    "grok-code-fast-1": {
+        "context": 256000,
+        "cost_input": 0.0,  # Free during promotional period
+        "cost_output": 0.0,
+        "supports_vision": False,
+        "supports_tools": True,
+        "reasoning_model": True,
+    },
+    # Grok 3 - Full model
+    "grok-3": {
+        "context": 128000,
+        "cost_input": 2.0,
+        "cost_output": 10.0,
+        "supports_vision": True,
+        "supports_tools": True,
+    },
+    "grok-3-fast": {
+        "context": 128000,
+        "cost_input": 1.0,
+        "cost_output": 5.0,
+        "supports_vision": True,
+        "supports_tools": True,
+    },
+    # Grok 3 Mini - Smaller, configurable reasoning
+    "grok-3-mini": {
+        "context": 128000,
+        "cost_input": 0.40,
+        "cost_output": 1.60,
+        "supports_vision": False,
+        "supports_tools": True,
+        "reasoning_model": True,  # Supports reasoning_effort parameter
+    },
+    "grok-3-mini-fast": {
+        "context": 128000,
+        "cost_input": 0.20,
+        "cost_output": 0.80,
+        "supports_vision": False,
+        "supports_tools": True,
+        "reasoning_model": True,
+    },
+    # Grok 2 - Legacy models
+    "grok-2-1212": {
+        "context": 131072,
+        "cost_input": 2.0,
+        "cost_output": 10.0,
+        "supports_vision": False,
+        "supports_tools": True,
+    },
+    "grok-2-vision": {
+        "context": 131072,
+        "cost_input": 2.0,
+        "cost_output": 10.0,
+        "supports_vision": True,
+        "supports_tools": True,
+    },
+    # Legacy aliases
     "grok-beta": {
         "context": 131072,
         "cost_input": 5.0,
@@ -1049,7 +1138,7 @@ INTERACTIVE_OPENAI_MODELS: Dict[str, Dict[str, Any]] = {
     },
     "gpt-4-turbo": {
         "display_name": "GPT-4 Turbo",
-        "description": "Legacy flagship, vision support",
+        "description": "Legacy flagship, tools support",
         "category": "Legacy Models",
     },
 }
@@ -1076,11 +1165,6 @@ INTERACTIVE_ANTHROPIC_MODELS: Dict[str, Dict[str, Any]] = {
         "description": "Proven stable, vision/tools",
         "category": "Claude 3.5 (Stable)",
     },
-    "claude-3-5-haiku-20241022": {
-        "display_name": "Claude 3.5 Haiku",
-        "description": "Budget option",
-        "category": "Claude 3.5 (Stable)",
-    },
 }
 
 # Google - 3 curated models
@@ -1140,16 +1224,46 @@ INTERACTIVE_GROQ_MODELS: Dict[str, Dict[str, Any]] = {
     },
 }
 
-# Grok (X.AI) - 1 curated model
+# Grok (X.AI) - 7 curated models
 INTERACTIVE_GROK_MODELS: Dict[str, Dict[str, Any]] = {
-    "grok-beta": {
-        "display_name": "Grok Beta",
-        "description": "X.AI flagship model",
-        "category": "Grok",
+    "grok-4-1-fast-reasoning": {
+        "display_name": "Grok 4.1 Fast (Reasoning)",
+        "description": "BEST VALUE - 2M context, agentic tools",
+        "category": "Grok 4.1 Fast (Latest)",
+    },
+    "grok-4-1-fast-non-reasoning": {
+        "display_name": "Grok 4.1 Fast (Non-Reasoning)",
+        "description": "FASTEST - instant responses, 2M context",
+        "category": "Grok 4.1 Fast (Latest)",
+    },
+    "grok-code-fast-1": {
+        "display_name": "Grok Code Fast",
+        "description": "FREE - specialized for coding",
+        "category": "Grok Specialized",
+    },
+    "grok-4": {
+        "display_name": "Grok 4",
+        "description": "Flagship reasoning model",
+        "category": "Grok 4 (Premium)",
+    },
+    "grok-3": {
+        "display_name": "Grok 3",
+        "description": "Stable production model",
+        "category": "Grok 3",
+    },
+    "grok-3-mini": {
+        "display_name": "Grok 3 Mini",
+        "description": "Configurable reasoning effort",
+        "category": "Grok 3",
+    },
+    "grok-2-1212": {
+        "display_name": "Grok 2",
+        "description": "Legacy stable model",
+        "category": "Grok 2 (Legacy)",
     },
 }
 
-# OpenRouter - 7 curated models (mix of free and paid)
+# OpenRouter - 11 curated models (mix of free and paid, multiple 1M context options)
 INTERACTIVE_OPENROUTER_MODELS: Dict[str, Dict[str, Any]] = {
     "anthropic/claude-sonnet-4-5": {
         "display_name": "Claude Sonnet 4.5",
@@ -1163,7 +1277,7 @@ INTERACTIVE_OPENROUTER_MODELS: Dict[str, Dict[str, Any]] = {
     },
     "google/gemini-2.5-flash": {
         "display_name": "Gemini 2.5 Flash",
-        "description": "Best value option",
+        "description": "BEST VALUE - 1M context, fast/cheap",
         "category": "Premium Models",
     },
     "meta-llama/llama-3.3-70b-instruct:free": {
@@ -1186,6 +1300,26 @@ INTERACTIVE_OPENROUTER_MODELS: Dict[str, Dict[str, Any]] = {
         "description": "European alternative",
         "category": "Premium Models",
     },
+    "anthropic/claude-opus-4-5": {
+        "display_name": "Claude Opus 4.5",
+        "description": "Premium quality, 1M context",
+        "category": "1M Context Models",
+    },
+    "google/gemini-2.5-pro": {
+        "display_name": "Gemini 2.5 Pro",
+        "description": "Best quality, 1M context",
+        "category": "1M Context Models",
+    },
+    "google/gemini-3": {
+        "display_name": "Gemini 3",
+        "description": "Latest Google, 1M context",
+        "category": "1M Context Models",
+    },
+    "google/gemini-2.0-flash-exp:free": {
+        "display_name": "Gemini 2.0 Flash Exp",
+        "description": "FREE - 1M context, vision/tools",
+        "category": "Free Models (1M Context)",
+    },
 }
 
 # Ollama - 3 curated models (local)

stratifyai/models.py

@@ -9,9 +9,44 @@ from typing import List, Literal, Optional
 class Message:
     """Standard message format for all providers (OpenAI-compatible)."""
     role: Literal["system", "user", "assistant"]
-    content: str
+    content: str  # Can be plain text or contain [IMAGE:mime_type]\nbase64_data format
     name: Optional[str] = None  # For multi-agent scenarios
     cache_control: Optional[dict] = None  # For providers that support prompt caching (Anthropic, OpenAI)
+    
+    def has_image(self) -> bool:
+        """Check if message contains image data."""
+        return "[IMAGE:" in self.content
+    
+    def parse_vision_content(self) -> tuple[Optional[str], Optional[tuple[str, str]]]:
+        """Parse content into text and image data.
+        
+        Returns:
+            (text_content, (mime_type, base64_data)) or (text_content, None) if no image
+        """
+        if not self.has_image():
+            return (self.content, None)
+        
+        # Split content by [IMAGE:...] marker
+        parts = self.content.split("[IMAGE:")
+        text_parts = []
+        image_data = None
+        
+        for i, part in enumerate(parts):
+            if i == 0:
+                # First part is text before image
+                if part.strip():
+                    text_parts.append(part.strip())
+            else:
+                # This part starts with mime_type]
+                if "]" in part:
+                    mime_type, rest = part.split("]", 1)
+                    # rest contains the base64 data (possibly with leading/trailing whitespace)
+                    base64_data = rest.strip()
+                    if base64_data:
+                        image_data = (mime_type.strip(), base64_data)
+        
+        text_content = "\n".join(text_parts).strip() if text_parts else None
+        return (text_content, image_data)
 
 
 @dataclass

stratifyai/providers/anthropic.py

@@ -94,7 +94,33 @@ class AnthropicProvider(BaseProvider):
             if msg.role == "system":
                 system_message = msg.content
             else:
-                message_dict = {"role": msg.role, "content": msg.content}
+                # Check if message contains image data
+                if msg.has_image():
+                    # Parse vision content
+                    text_content, image_data = msg.parse_vision_content()
+                    
+                    # Build vision message content array
+                    content_parts = []
+                    if text_content:
+                        content_parts.append({"type": "text", "text": text_content})
+                    
+                    if image_data:
+                        mime_type, base64_data = image_data
+                        # Anthropic expects base64 with source
+                        content_parts.append({
+                            "type": "image",
+                            "source": {
+                                "type": "base64",
+                                "media_type": mime_type,
+                                "data": base64_data
+                            }
+                        })
+                    
+                    message_dict = {"role": msg.role, "content": content_parts}
+                else:
+                    # Regular text message
+                    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
@@ -138,8 +164,16 @@ class AnthropicProvider(BaseProvider):
             # Normalize and return
             return self._normalize_response(raw_response.model_dump())
         except Exception as e:
+            error_str = str(e)
+            # Check for vision-related errors
+            if "image" in error_str.lower() and ("not supported" in error_str.lower() or "invalid" in error_str.lower()):
+                raise ProviderAPIError(
+                    f"Vision not supported: The model '{request.model}' cannot process images. "
+                    f"Please use a vision-capable Claude model like 'claude-sonnet-4-5' or 'claude-opus-4-5'.",
+                    self.provider_name
+                )
             raise ProviderAPIError(
-                f"Chat completion failed: {str(e)}",
+                f"Chat completion failed: {error_str}",
                 self.provider_name
             )
     
@@ -170,7 +204,7 @@ class AnthropicProvider(BaseProvider):
             constraints.get("max_temperature", 1.0)
         )
         
-        # Convert messages to Anthropic format
+        # Convert messages to Anthropic format with vision support
         system_message = None
         messages = []
         
@@ -178,7 +212,25 @@ class AnthropicProvider(BaseProvider):
             if msg.role == "system":
                 system_message = msg.content
             else:
-                messages.append({"role": msg.role, "content": msg.content})
+                if msg.has_image():
+                    # Parse and format vision content
+                    text_content, image_data = msg.parse_vision_content()
+                    content_parts = []
+                    if text_content:
+                        content_parts.append({"type": "text", "text": text_content})
+                    if image_data:
+                        mime_type, base64_data = image_data
+                        content_parts.append({
+                            "type": "image",
+                            "source": {
+                                "type": "base64",
+                                "media_type": mime_type,
+                                "data": base64_data
+                            }
+                        })
+                    messages.append({"role": msg.role, "content": content_parts})
+                else:
+                    messages.append({"role": msg.role, "content": msg.content})
         
         # Build request parameters
         anthropic_params = {
@@ -196,8 +248,16 @@ class AnthropicProvider(BaseProvider):
                 async for chunk in stream.text_stream:
                     yield self._normalize_stream_chunk(chunk)
         except Exception as e:
+            error_str = str(e)
+            # Check for vision-related errors
+            if "image" in error_str.lower() and ("not supported" in error_str.lower() or "invalid" in error_str.lower()):
+                raise ProviderAPIError(
+                    f"Vision not supported: The model '{request.model}' cannot process images. "
+                    f"Please use a vision-capable Claude model like 'claude-sonnet-4-5' or 'claude-opus-4-5'.",
+                    self.provider_name
+                )
             raise ProviderAPIError(
-                f"Streaming chat completion failed: {str(e)}",
+                f"Streaming chat completion failed: {error_str}",
                 self.provider_name
             )
     

stratifyai/providers/bedrock.py

@@ -161,6 +161,21 @@ class BedrockProvider(BaseProvider):
         except ClientError as e:
             error_code = e.response["Error"]["Code"]
             error_message = e.response["Error"]["Message"]
+            
+            # Parse and provide user-friendly error messages
+            if error_code == "ValidationException":
+                # Extract specific validation issues
+                if "is not less or equal to" in error_message and "/p:" in error_message:
+                    friendly_msg = "Model configuration error: top_p parameter exceeds maximum allowed value for this model."
+                elif "is not a valid enum value" in error_message and "role" in error_message:
+                    friendly_msg = "Model configuration error: Invalid message role format for this model."
+                else:
+                    friendly_msg = f"Request validation failed: {error_message}"
+                raise ProviderAPIError(
+                    f"[bedrock] {friendly_msg}",
+                    self.provider_name
+                )
+            
             raise ProviderAPIError(
                 f"Bedrock API error ({error_code}): {error_message}",
                 self.provider_name
@@ -223,6 +238,20 @@ class BedrockProvider(BaseProvider):
         except ClientError as e:
             error_code = e.response["Error"]["Code"]
             error_message = e.response["Error"]["Message"]
+            
+            # Parse and provide user-friendly error messages
+            if error_code == "ValidationException":
+                if "is not less or equal to" in error_message and "/p:" in error_message:
+                    friendly_msg = "Model configuration error: top_p parameter exceeds maximum allowed value for this model."
+                elif "is not a valid enum value" in error_message and "role" in error_message:
+                    friendly_msg = "Model configuration error: Invalid message role format for this model."
+                else:
+                    friendly_msg = f"Request validation failed: {error_message}"
+                raise ProviderAPIError(
+                    f"[bedrock] {friendly_msg}",
+                    self.provider_name
+                )
+            
             raise ProviderAPIError(
                 f"Bedrock streaming error ({error_code}): {error_message}",
                 self.provider_name
@@ -292,7 +321,29 @@ class BedrockProvider(BaseProvider):
             if msg.role == "system":
                 system_message = msg.content
             else:
-                messages.append({"role": msg.role, "content": msg.content})
+                # Check if message contains an image
+                if msg.has_image():
+                    # Parse vision content
+                    text_content, (mime_type, base64_data) = msg.parse_vision_content()
+                    
+                    # Build content array for vision (Anthropic format)
+                    content_parts = []
+                    if text_content:
+                        content_parts.append({"type": "text", "text": text_content})
+                    
+                    # Add image in Anthropic format
+                    content_parts.append({
+                        "type": "image",
+                        "source": {
+                            "type": "base64",
+                            "media_type": mime_type,
+                            "data": base64_data
+                        }
+                    })
+                    
+                    messages.append({"role": msg.role, "content": content_parts})
+                else:
+                    messages.append({"role": msg.role, "content": msg.content})
         
         body = {
             "anthropic_version": "bedrock-2023-05-31",
@@ -338,17 +389,33 @@ class BedrockProvider(BaseProvider):
     
     def _build_cohere_request(self, request: ChatRequest) -> dict:
         """Build request for Cohere models."""
-        # Cohere uses a message-based format similar to OpenAI
-        messages = []
-        for msg in request.messages:
-            messages.append({"role": msg.role, "message": msg.content})
+        # Cohere Bedrock uses USER/CHATBOT roles and requires specific format
+        # Extract user message (last message should be from user)
+        user_message = ""
+        chat_history = []
+        
+        for i, msg in enumerate(request.messages):
+            # Skip system messages - Cohere handles them differently
+            if msg.role == "system":
+                continue
+            
+            # Last user message becomes the main message
+            if i == len(request.messages) - 1 and msg.role == "user":
+                user_message = msg.content
+            else:
+                # Map role names to Cohere's expected format
+                cohere_role = "USER" if msg.role == "user" else "CHATBOT"
+                chat_history.append({"role": cohere_role, "message": msg.content})
+        
+        # Clamp top_p to Cohere's maximum of 0.99
+        top_p = min(request.top_p, 0.99)
         
         return {
-            "message": messages[-1]["message"] if messages else "",
-            "chat_history": messages[:-1] if len(messages) > 1 else [],
+            "message": user_message,
+            "chat_history": chat_history,
             "max_tokens": request.max_tokens or 2048,
             "temperature": request.temperature,
-            "p": request.top_p,
+            "p": top_p,
         }
     
     def _build_nova_request(self, request: ChatRequest) -> dict:
@@ -361,7 +428,27 @@ class BedrockProvider(BaseProvider):
             if msg.role == "system":
                 system_message = msg.content
             else:
-                messages.append({"role": msg.role, "content": [{"text": msg.content}]})
+                # Check if message contains an image
+                if msg.has_image():
+                    # Parse vision content
+                    text_content, (mime_type, base64_data) = msg.parse_vision_content()
+                    
+                    # Build content array for vision (Nova format)
+                    content_parts = []
+                    if text_content:
+                        content_parts.append({"text": text_content})
+                    
+                    # Add image in Nova format
+                    content_parts.append({
+                        "image": {
+                            "format": mime_type.split("/")[1] if "/" in mime_type else "png",
+                            "source": {"bytes": base64_data}
+                        }
+                    })
+                    
+                    messages.append({"role": msg.role, "content": content_parts})
+                else:
+                    messages.append({"role": msg.role, "content": [{"text": msg.content}]})
         
         body = {
             "messages": messages,

stratifyai/providers/grok.py

@@ -20,13 +20,14 @@ class GrokProvider(OpenAICompatibleProvider):
         Initialize Grok provider.
         
         Args:
-            api_key: Grok API key (defaults to GROK_API_KEY env var)
+            api_key: Grok API key (defaults to XAI_API_KEY or 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")
+        # Support both XAI_API_KEY (official) and GROK_API_KEY (legacy) for backward compatibility
+        api_key = api_key or os.getenv("XAI_API_KEY") or os.getenv("GROK_API_KEY")
         if not api_key:
             raise AuthenticationError("grok")
         

stratifyai/providers/openai.py

@@ -87,7 +87,30 @@ class OpenAIProvider(BaseProvider):
         # Build OpenAI-specific request parameters
         messages = []
         for msg in request.messages:
-            message_dict = {"role": msg.role, "content": msg.content}
+            # Check if message contains image data
+            if msg.has_image():
+                # Parse vision content
+                text_content, image_data = msg.parse_vision_content()
+                
+                # Build vision message content array
+                content_parts = []
+                if text_content:
+                    content_parts.append({"type": "text", "text": text_content})
+                
+                if image_data:
+                    mime_type, base64_data = image_data
+                    # OpenAI expects data URL format
+                    image_url = f"data:{mime_type};base64,{base64_data}"
+                    content_parts.append({
+                        "type": "image_url",
+                        "image_url": {"url": image_url}
+                    })
+                
+                message_dict = {"role": msg.role, "content": content_parts}
+            else:
+                # Regular text message
+                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
@@ -143,8 +166,16 @@ class OpenAIProvider(BaseProvider):
             # Normalize and return
             return self._normalize_response(raw_response.model_dump())
         except Exception as e:
+            error_str = str(e)
+            # Check for vision-related errors
+            if "image_url is only supported by certain models" in error_str or "Invalid content type" in error_str:
+                raise ProviderAPIError(
+                    f"Vision not supported: The model '{request.model}' cannot process images. "
+                    f"Please use a vision-capable model like 'gpt-4o' or 'gpt-4o-mini'.",
+                    self.provider_name
+                )
             raise ProviderAPIError(
-                f"Chat completion failed: {str(e)}",
+                f"Chat completion failed: {error_str}",
                 self.provider_name
             )
     
@@ -167,13 +198,29 @@ class OpenAIProvider(BaseProvider):
         if not self.validate_model(request.model):
             raise InvalidModelError(request.model, self.provider_name)
         
-        # Build request parameters
+        # Build request parameters with vision support
+        messages = []
+        for msg in request.messages:
+            if msg.has_image():
+                # Parse and format vision content
+                text_content, image_data = msg.parse_vision_content()
+                content_parts = []
+                if text_content:
+                    content_parts.append({"type": "text", "text": text_content})
+                if image_data:
+                    mime_type, base64_data = image_data
+                    image_url = f"data:{mime_type};base64,{base64_data}"
+                    content_parts.append({
+                        "type": "image_url",
+                        "image_url": {"url": image_url}
+                    })
+                messages.append({"role": msg.role, "content": content_parts})
+            else:
+                messages.append({"role": msg.role, "content": msg.content})
+        
         openai_params = {
             "model": request.model,
-            "messages": [
-                {"role": msg.role, "content": msg.content}
-                for msg in request.messages
-            ],
+            "messages": messages,
             "stream": True,
         }
         
@@ -207,8 +254,16 @@ class OpenAIProvider(BaseProvider):
                 if chunk.choices and chunk.choices[0].delta.content:
                     yield self._normalize_stream_chunk(chunk_dict)
         except Exception as e:
+            error_str = str(e)
+            # Check for vision-related errors
+            if "image_url is only supported by certain models" in error_str or "Invalid content type" in error_str:
+                raise ProviderAPIError(
+                    f"Vision not supported: The model '{request.model}' cannot process images. "
+                    f"Please use a vision-capable model like 'gpt-4o' or 'gpt-4o-mini'.",
+                    self.provider_name
+                )
             raise ProviderAPIError(
-                f"Streaming chat completion failed: {str(e)}",
+                f"Streaming chat completion failed: {error_str}",
                 self.provider_name
             )