mito-ai 0.1.57__py3-none-any.whl → 0.1.59__py3-none-any.whl

mito_ai/tests/test_model_utils.py

@@ -0,0 +1,362 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+from unittest.mock import patch, MagicMock
+from mito_ai.utils.model_utils import (
+    get_available_models,
+    get_fast_model_for_selected_model,
+    STANDARD_MODELS,
+    ANTHROPIC_MODEL_ORDER,
+    OPENAI_MODEL_ORDER,
+    GEMINI_MODEL_ORDER,
+)
+
+
+class TestGetAvailableModels:
+    """Tests for get_available_models() function."""
+    
+    @patch('mito_ai.utils.model_utils.is_enterprise')
+    @patch('mito_ai.utils.model_utils.constants')
+    def test_returns_litellm_models_when_enterprise_and_configured(self, mock_constants, mock_is_enterprise):
+        """Test that LiteLLM models are returned when enterprise mode is enabled and LiteLLM is configured."""
+        mock_is_enterprise.return_value = True
+        mock_constants.LITELLM_BASE_URL = "https://litellm-server.com"
+        mock_constants.LITELLM_MODELS = ["litellm/openai/gpt-4o", "litellm/anthropic/claude-3-5-sonnet"]
+        
+        result = get_available_models()
+        
+        assert result == ["litellm/openai/gpt-4o", "litellm/anthropic/claude-3-5-sonnet"]
+    
+    @patch('mito_ai.utils.model_utils.is_enterprise')
+    @patch('mito_ai.utils.model_utils.constants')
+    def test_returns_standard_models_when_not_enterprise(self, mock_constants, mock_is_enterprise):
+        """Test that standard models are returned when enterprise mode is not enabled."""
+        mock_is_enterprise.return_value = False
+        
+        result = get_available_models()
+        
+        assert result == STANDARD_MODELS
+    
+    @patch('mito_ai.utils.model_utils.is_enterprise')
+    @patch('mito_ai.utils.model_utils.constants')
+    def test_returns_standard_models_when_enterprise_but_no_litellm(self, mock_constants, mock_is_enterprise):
+        """Test that standard models are returned when enterprise mode is enabled but LiteLLM is not configured."""
+        mock_is_enterprise.return_value = True
+        mock_constants.LITELLM_BASE_URL = None
+        mock_constants.LITELLM_MODELS = []
+        
+        result = get_available_models()
+        
+        assert result == STANDARD_MODELS
+    
+    @patch('mito_ai.utils.model_utils.is_enterprise')
+    @patch('mito_ai.utils.model_utils.constants')
+    def test_returns_standard_models_when_enterprise_but_no_base_url(self, mock_constants, mock_is_enterprise):
+        """Test that standard models are returned when enterprise mode is enabled but LITELLM_BASE_URL is not set."""
+        mock_is_enterprise.return_value = True
+        mock_constants.LITELLM_BASE_URL = None
+        mock_constants.LITELLM_MODELS = ["litellm/openai/gpt-4o"]
+        
+        result = get_available_models()
+        
+        assert result == STANDARD_MODELS
+    
+    @patch('mito_ai.utils.model_utils.is_enterprise')
+    @patch('mito_ai.utils.model_utils.constants')
+    def test_returns_standard_models_when_enterprise_but_no_models(self, mock_constants, mock_is_enterprise):
+        """Test that standard models are returned when enterprise mode is enabled but LITELLM_MODELS is empty."""
+        mock_is_enterprise.return_value = True
+        mock_constants.LITELLM_BASE_URL = "https://litellm-server.com"
+        mock_constants.LITELLM_MODELS = []
+        
+        result = get_available_models()
+        
+        assert result == STANDARD_MODELS
+    
+    @patch('mito_ai.utils.model_utils.is_abacus_configured')
+    @patch('mito_ai.utils.model_utils.is_enterprise')
+    @patch('mito_ai.utils.model_utils.constants')
+    def test_returns_abacus_models_when_configured(self, mock_constants, mock_is_enterprise, mock_is_abacus_configured):
+        """Test that Abacus models are returned when Abacus is configured (highest priority)."""
+        mock_is_abacus_configured.return_value = True
+        mock_is_enterprise.return_value = True
+        mock_constants.ABACUS_BASE_URL = "https://routellm.abacus.ai/v1"
+        mock_constants.ABACUS_MODELS = ["Abacus/gpt-4.1", "Abacus/claude-haiku-4-5-20251001"]
+        
+        result = get_available_models()
+        
+        assert result == ["Abacus/gpt-4.1", "Abacus/claude-haiku-4-5-20251001"]
+
+
+class TestGetFastModelForSelectedModel:
+    """Tests for get_fast_model_for_selected_model() function."""
+    
+    def test_anthropic_sonnet_returns_haiku(self):
+        """Test that Claude Sonnet returns Claude Haiku (fastest Anthropic model)."""
+        result = get_fast_model_for_selected_model("claude-sonnet-4-5-20250929")
+        assert result == "claude-haiku-4-5-20251001"
+    
+    def test_anthropic_haiku_returns_haiku(self):
+        """Test that Claude Haiku returns itself (already fastest)."""
+        result = get_fast_model_for_selected_model("claude-haiku-4-5-20251001")
+        assert result == "claude-haiku-4-5-20251001"
+    
+    def test_openai_gpt_4_1_returns_gpt_4_1(self):
+        """Test that GPT 4.1 returns itself (already fastest)."""
+        result = get_fast_model_for_selected_model("gpt-4.1")
+        assert result == "gpt-4.1"
+    
+    def test_openai_gpt_5_2_returns_gpt_4_1(self):
+        """Test that GPT 5.2 returns GPT 4.1 (fastest OpenAI model)."""
+        result = get_fast_model_for_selected_model("gpt-5.2")
+        assert result == "gpt-4.1"
+    
+    def test_gemini_pro_returns_flash(self):
+        """Test that Gemini Pro returns Gemini Flash (fastest Gemini model)."""
+        result = get_fast_model_for_selected_model("gemini-3-pro-preview")
+        assert result == "gemini-3-flash-preview"
+    
+    def test_gemini_flash_returns_flash(self):
+        """Test that Gemini Flash returns itself (already fastest)."""
+        result = get_fast_model_for_selected_model("gemini-3-flash-preview")
+        assert result == "gemini-3-flash-preview"
+    
+    @patch('mito_ai.utils.model_utils.get_available_models')
+    @pytest.mark.parametrize(
+        "selected_model,available_models,expected_result",
+        [
+            # Test case 1: LiteLLM OpenAI model returns fastest overall
+            (
+                "litellm/openai/gpt-5.2",
+                ["litellm/openai/gpt-4.1", "litellm/openai/gpt-5.2", "litellm/anthropic/claude-sonnet-4-5-20250929"],
+                "litellm/openai/gpt-4.1",
+            ),
+            # Test case 2: LiteLLM Anthropic model returns fastest overall
+            (
+                "litellm/anthropic/claude-sonnet-4-5-20250929",
+                ["litellm/openai/gpt-4.1", "litellm/anthropic/claude-sonnet-4-5-20250929", "litellm/anthropic/claude-haiku-4-5-20251001"],
+                "litellm/openai/gpt-4.1",
+            ),
+            # Test case 3: LiteLLM Google model returns fastest overall
+            (
+                "litellm/google/gemini-3-pro-preview",
+                ["litellm/google/gemini-3-pro-preview", "litellm/google/gemini-3-flash-preview"],
+                "litellm/google/gemini-3-flash-preview",
+            ),
+            # Test case 4: Unknown LiteLLM model returns fastest known
+            (
+                "unknown/provider/model",
+                ["litellm/openai/gpt-4.1", "unknown/provider/model"],
+                "litellm/openai/gpt-4.1",
+            ),
+            # Test case 5: Single LiteLLM model returns itself
+            (
+                "litellm/openai/gpt-4o",
+                ["litellm/openai/gpt-4o"],
+                "litellm/openai/gpt-4o",
+            ),
+            # Test case 6: Cross-provider comparison - OpenAI is faster
+            (
+                "litellm/anthropic/claude-sonnet-4-5-20250929",
+                [
+                    "litellm/openai/gpt-4.1",  # Index 0 in OPENAI_MODEL_ORDER
+                    "litellm/anthropic/claude-sonnet-4-5-20250929",  # Index 1 in ANTHROPIC_MODEL_ORDER
+                ],
+                "litellm/openai/gpt-4.1",
+            ),
+            # Test case 7: Cross-provider comparison - Anthropic is faster
+            (
+                "litellm/openai/gpt-5.2",
+                [
+                    "litellm/openai/gpt-5.2",  # Index 1 in OPENAI_MODEL_ORDER
+                    "litellm/anthropic/claude-haiku-4-5-20251001",  # Index 0 in ANTHROPIC_MODEL_ORDER
+                ],
+                "litellm/anthropic/claude-haiku-4-5-20251001",
+            ),
+        ],
+        ids=[
+            "litellm_openai_model_returns_fastest_overall",
+            "litellm_anthropic_model_returns_fastest_overall",
+            "litellm_google_model_returns_fastest_overall",
+            "litellm_unknown_model_returns_fastest_known",
+            "litellm_single_model_returns_itself",
+            "litellm_cross_provider_comparison_openai_faster",
+            "litellm_returns_fastest_when_anthropic_is_faster",
+        ]
+    )
+    def test_litellm_model_returns_fastest(
+        self,
+        mock_get_available_models,
+        selected_model,
+        available_models,
+        expected_result,
+    ):
+        """Test that LiteLLM models return fastest model from all available models."""
+        mock_get_available_models.return_value = available_models
+        
+        result = get_fast_model_for_selected_model(selected_model)
+        
+        assert result == expected_result
+    
+    def test_unknown_standard_model_returns_itself(self):
+        """Test that unknown standard model returns itself."""
+        result = get_fast_model_for_selected_model("unknown-model")
+        assert result == "unknown-model"
+    
+    def test_claude_model_not_in_order_returns_fastest_anthropic(self):
+        """Test that a Claude model not in ANTHROPIC_MODEL_ORDER still returns fastest Anthropic model."""
+        # Test with a Claude model that isn't in the order list
+        result = get_fast_model_for_selected_model("claude-3-opus-20240229")
+        # Should return fastest Anthropic model (claude-haiku-4-5-20251001)
+        assert result == "claude-haiku-4-5-20251001"
+        assert result.startswith("claude")
+    
+    def test_gpt_model_not_in_order_returns_fastest_openai(self):
+        """Test that a GPT model not in OPENAI_MODEL_ORDER still returns fastest OpenAI model."""
+        # Test with a GPT model that isn't in the order list
+        result = get_fast_model_for_selected_model("gpt-4o-mini")
+        # Should return fastest OpenAI model (gpt-4.1)
+        assert result == "gpt-4.1"
+        assert result.startswith("gpt")
+    
+    def test_gemini_model_not_in_order_returns_fastest_gemini(self):
+        """Test that a Gemini model not in GEMINI_MODEL_ORDER still returns fastest Gemini model."""
+        # Test with a Gemini model that isn't in the order list
+        result = get_fast_model_for_selected_model("gemini-1.5-pro")
+        # Should return fastest Gemini model (gemini-3-flash-preview)
+        assert result == "gemini-3-flash-preview"
+        assert result.startswith("gemini")
+    
+    def test_claude_model_variations_return_same_provider(self):
+        """Test that various Claude model name variations return Anthropic models."""
+        test_cases = [
+            "claude-3-5-sonnet",
+            "claude-3-opus",
+            "claude-instant",
+            "claude-v2",
+        ]
+        for model in test_cases:
+            result = get_fast_model_for_selected_model(model)
+            # Should always return an Anthropic model (starts with "claude")
+            assert result.startswith("claude"), f"Model {model} should return Anthropic model, got {result}"
+            # Should return the fastest Anthropic model
+            assert result == "claude-haiku-4-5-20251001", f"Model {model} should return fastest Anthropic model"
+    
+    def test_gpt_model_variations_return_same_provider(self):
+        """Test that various GPT model name variations return OpenAI models."""
+        test_cases = [
+            "gpt-4o",
+            "gpt-4-turbo",
+            "gpt-3.5-turbo",
+            "gpt-4o-mini",
+        ]
+        for model in test_cases:
+            result = get_fast_model_for_selected_model(model)
+            # Should always return an OpenAI model (starts with "gpt")
+            assert result.startswith("gpt"), f"Model {model} should return OpenAI model, got {result}"
+            # Should return the fastest OpenAI model
+            assert result == "gpt-4.1", f"Model {model} should return fastest OpenAI model"
+    
+    def test_gemini_model_variations_return_same_provider(self):
+        """Test that various Gemini model name variations return Gemini models."""
+        test_cases = [
+            "gemini-1.5-pro",
+            "gemini-1.5-flash",
+            "gemini-pro",
+            "gemini-ultra",
+        ]
+        for model in test_cases:
+            result = get_fast_model_for_selected_model(model)
+            # Should always return a Gemini model (starts with "gemini")
+            assert result.startswith("gemini"), f"Model {model} should return Gemini model, got {result}"
+            # Should return the fastest Gemini model
+            assert result == "gemini-3-flash-preview", f"Model {model} should return fastest Gemini model"
+    
+    def test_case_insensitive_provider_matching(self):
+        """Test that provider matching is case-insensitive."""
+        test_cases = [
+            ("CLAUDE-sonnet-4-5-20250929", "claude-haiku-4-5-20251001"),
+            ("GPT-4.1", "gpt-4.1"),
+            ("GEMINI-3-flash-preview", "gemini-3-flash-preview"),
+        ]
+        for model, expected in test_cases:
+            result = get_fast_model_for_selected_model(model)
+            assert result == expected, f"Case-insensitive matching failed for {model}"
+    
+    @patch('mito_ai.utils.model_utils.get_available_models')
+    @pytest.mark.parametrize(
+        "selected_model,available_models,expected_result",
+        [
+            # Test case 1: Abacus GPT model returns fastest overall
+            (
+                "Abacus/gpt-5.2",
+                ["Abacus/gpt-4.1", "Abacus/gpt-5.2", "Abacus/claude-sonnet-4-5-20250929"],
+                "Abacus/gpt-4.1",
+            ),
+            # Test case 2: Abacus Claude model returns fastest overall
+            (
+                "Abacus/claude-sonnet-4-5-20250929",
+                ["Abacus/gpt-4.1", "Abacus/claude-sonnet-4-5-20250929", "Abacus/claude-haiku-4-5-20251001"],
+                "Abacus/gpt-4.1",
+            ),
+            # Test case 3: Abacus Gemini model returns fastest overall
+            (
+                "Abacus/gemini-3-pro-preview",
+                ["Abacus/gemini-3-pro-preview", "Abacus/gemini-3-flash-preview"],
+                "Abacus/gemini-3-flash-preview",
+            ),
+            # Test case 4: Unknown Abacus model returns fastest known
+            (
+                "Abacus/unknown-model",
+                ["Abacus/gpt-4.1", "Abacus/unknown-model"],
+                "Abacus/gpt-4.1",
+            ),
+            # Test case 5: Single Abacus model returns itself
+            (
+                "Abacus/gpt-4.1",
+                ["Abacus/gpt-4.1"],
+                "Abacus/gpt-4.1",
+            ),
+            # Test case 6: Cross-provider comparison - OpenAI is faster
+            (
+                "Abacus/claude-sonnet-4-5-20250929",
+                [
+                    "Abacus/gpt-4.1",  # Index 0 in OPENAI_MODEL_ORDER
+                    "Abacus/claude-sonnet-4-5-20250929",  # Index 1 in ANTHROPIC_MODEL_ORDER
+                ],
+                "Abacus/gpt-4.1",
+            ),
+            # Test case 7: Cross-provider comparison - Anthropic is faster
+            (
+                "Abacus/gpt-5.2",
+                [
+                    "Abacus/gpt-5.2",  # Index 1 in OPENAI_MODEL_ORDER
+                    "Abacus/claude-haiku-4-5-20251001",  # Index 0 in ANTHROPIC_MODEL_ORDER
+                ],
+                "Abacus/claude-haiku-4-5-20251001",
+            ),
+        ],
+        ids=[
+            "abacus_gpt_model_returns_fastest_overall",
+            "abacus_anthropic_model_returns_fastest_overall",
+            "abacus_google_model_returns_fastest_overall",
+            "abacus_unknown_model_returns_fastest_known",
+            "abacus_single_model_returns_itself",
+            "abacus_cross_provider_comparison_openai_faster",
+            "abacus_returns_fastest_when_anthropic_is_faster",
+        ]
+    )
+    def test_abacus_model_returns_fastest(
+        self,
+        mock_get_available_models,
+        selected_model,
+        available_models,
+        expected_result,
+    ):
+        """Test that Abacus models return fastest model from all available models."""
+        mock_get_available_models.return_value = available_models
+        
+        result = get_fast_model_for_selected_model(selected_model)
+        
+        assert result == expected_result

mito_ai/utils/anthropic_utils.py

@@ -5,7 +5,6 @@ import anthropic
 from typing import Any, Dict, List, Optional, Union, AsyncGenerator, Tuple, Callable
 from anthropic.types import MessageParam, TextBlockParam, ToolUnionParam
 from mito_ai.utils.mito_server_utils import get_response_from_mito_server, stream_response_from_mito_server
-from mito_ai.utils.provider_utils import does_message_require_fast_model
 from mito_ai.completions.models import AgentResponse, MessageType, ResponseFormatInfo, CompletionReply, CompletionStreamChunk
 from mito_ai.utils.schema import UJ_STATIC_USER_ID, UJ_USER_EMAIL
 from mito_ai.utils.db import get_user_field
@@ -20,6 +19,7 @@ max_retries = 1
 
 FAST_ANTHROPIC_MODEL = "claude-haiku-4-5-20251001" # This should be in sync with ModelSelector.tsx
 LARGE_CONTEXT_MODEL = "claude-sonnet-4-5-20250929" # This should be in sync with ModelSelector.tsx
+EXTENDED_CONTEXT_BETA = "context-1m-2025-08-07" # Beta feature for extended context window support
 
 def does_message_exceed_max_tokens(system: Union[str, List[TextBlockParam], anthropic.Omit], messages: List[MessageParam]) -> bool:
     token_estimation = get_rough_token_estimatation_anthropic(system, messages)
@@ -36,10 +36,6 @@ def select_correct_model(default_model: str, message_type: MessageType, system:
         # but not haiku models
         return LARGE_CONTEXT_MODEL
     
-    message_requires_fast_model = does_message_require_fast_model(message_type)
-    if message_requires_fast_model:
-        return FAST_ANTHROPIC_MODEL
-    
     return default_model    
 
 def _prepare_anthropic_request_data_and_headers(
@@ -66,7 +62,7 @@ def _prepare_anthropic_request_data_and_headers(
         "max_tokens": max_tokens,
         "temperature": temperature,
         "messages": messages,
-        "betas": ["context-1m-2025-08-07"]
+        "betas": [EXTENDED_CONTEXT_BETA]
     }
     
     # Add system to inner_data only if it is not anthropic.Omit
@@ -173,6 +169,12 @@ def get_anthropic_completion_function_params(
         "messages": messages,
         "system": system,
     }
+    
+    # Enable extended context beta when using LARGE_CONTEXT_MODEL
+    # This is required for messages exceeding the standard context limit
+    if model == LARGE_CONTEXT_MODEL:
+        provider_data["betas"] = [EXTENDED_CONTEXT_BETA]
+    
     if tools:
         provider_data["tools"] = tools
     if response_format_info and response_format_info.name == "agent_response":

mito_ai/utils/gemini_utils.py

@@ -8,7 +8,6 @@ from typing import Any, Dict, List, Optional, Callable, Union, AsyncGenerator, T
 from mito_ai.utils.mito_server_utils import get_response_from_mito_server, stream_response_from_mito_server
 from mito_ai.completions.models import AgentResponse, CompletionReply, CompletionStreamChunk, CompletionItem, MessageType
 from mito_ai.constants import MITO_GEMINI_URL
-from mito_ai.utils.provider_utils import does_message_require_fast_model
 from mito_ai.utils.utils import _create_http_client
 
 timeout = 30
@@ -114,8 +113,6 @@ def get_gemini_completion_function_params(
     Build the provider_data dict for Gemini completions, mirroring the OpenAI/Anthropic approach.
     Only includes fields needed for the Gemini API.
     """
-    message_requires_fast_model = does_message_require_fast_model(message_type)
-    model = FAST_GEMINI_MODEL if message_requires_fast_model else model
     
     provider_data: Dict[str, Any] = {
         "model": model,

mito_ai/utils/litellm_utils.py

@@ -0,0 +1,84 @@
+#!/usr/bin/env python
+# coding: utf-8
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the The Mito Enterprise license.
+
+from typing import Dict, Any, List, Optional
+from openai.types.chat import ChatCompletionMessageParam
+import copy
+
+from mito_ai import constants
+from mito_ai.completions.models import ResponseFormatInfo
+from mito_ai.utils.version_utils import is_enterprise
+
+def is_litellm_configured() -> bool:
+    """
+    Check if LiteLLM is configured for system use.
+    
+    Per enterprise documentation, LITELLM_API_KEY is user-controlled and optional
+    for system configuration. This function only checks system-level configuration
+    (BASE_URL and MODELS), not user-specific API keys.
+    """
+    return all([constants.LITELLM_BASE_URL, constants.LITELLM_MODELS, is_enterprise()])
+
+def get_litellm_completion_function_params(
+    model: str,
+    messages: List[ChatCompletionMessageParam],
+    api_key: Optional[str],
+    api_base: str,
+    timeout: int,
+    stream: bool,
+    response_format_info: Optional[ResponseFormatInfo] = None,
+) -> Dict[str, Any]:
+    """
+    Prepare parameters for LiteLLM completion requests.
+    
+    Args:
+        model: Model name with provider prefix (e.g., "openai/gpt-4o")
+        messages: List of chat messages
+        api_key: Optional API key for authentication
+        api_base: Base URL for the LiteLLM server
+        timeout: Request timeout in seconds
+        stream: Whether to stream the response
+        response_format_info: Optional response format specification
+        
+    Returns:
+        Dictionary of parameters ready to be passed to litellm.acompletion()
+    """
+    params: Dict[str, Any] = {
+        "model": model,
+        "messages": messages,
+        "api_key": api_key,
+        "api_base": api_base,
+        "timeout": timeout,
+        "stream": stream,
+    }
+    
+    # Handle response format if specified
+    if response_format_info:
+        # LiteLLM supports response_format for structured outputs
+        if hasattr(response_format_info.format, 'model_json_schema'):
+            # Pydantic model - get JSON schema
+            # Make a deep copy to avoid mutating the original schema
+            schema = copy.deepcopy(response_format_info.format.model_json_schema())
+            
+            # Add additionalProperties: False to the top-level schema
+            # This is required by OpenAI's JSON schema mode
+            schema["additionalProperties"] = False
+            
+            # Nested object definitions in $defs need to have additionalProperties set to False also
+            if "$defs" in schema:
+                for def_name, def_schema in schema["$defs"].items():
+                    if def_schema.get("type") == "object":
+                        def_schema["additionalProperties"] = False
+            
+            params["response_format"] = {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": response_format_info.name,
+                    "schema": schema,
+                    "strict": True
+                }
+            }
+    
+    return params