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

mito_ai/tests/test_model_utils.py

@@ -0,0 +1,271 @@
+# 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 = ["openai/gpt-4o", "anthropic/claude-3-5-sonnet"]
+        
+        result = get_available_models()
+        
+        assert result == ["openai/gpt-4o", "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 = ["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
+
+
+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
+            (
+                "openai/gpt-5.2",
+                ["openai/gpt-4.1", "openai/gpt-5.2", "anthropic/claude-sonnet-4-5-20250929"],
+                "openai/gpt-4.1",
+            ),
+            # Test case 2: LiteLLM Anthropic model returns fastest overall
+            (
+                "anthropic/claude-sonnet-4-5-20250929",
+                ["openai/gpt-4.1", "anthropic/claude-sonnet-4-5-20250929", "anthropic/claude-haiku-4-5-20251001"],
+                "openai/gpt-4.1",
+            ),
+            # Test case 3: LiteLLM Google model returns fastest overall
+            (
+                "google/gemini-3-pro-preview",
+                ["google/gemini-3-pro-preview", "google/gemini-3-flash-preview"],
+                "google/gemini-3-flash-preview",
+            ),
+            # Test case 4: Unknown LiteLLM model returns fastest known
+            (
+                "unknown/provider/model",
+                ["openai/gpt-4.1", "unknown/provider/model"],
+                "openai/gpt-4.1",
+            ),
+            # Test case 5: Single LiteLLM model returns itself
+            (
+                "openai/gpt-4o",
+                ["openai/gpt-4o"],
+                "openai/gpt-4o",
+            ),
+            # Test case 6: Cross-provider comparison - OpenAI is faster
+            (
+                "anthropic/claude-sonnet-4-5-20250929",
+                [
+                    "openai/gpt-4.1",  # Index 0 in OPENAI_MODEL_ORDER
+                    "anthropic/claude-sonnet-4-5-20250929",  # Index 1 in ANTHROPIC_MODEL_ORDER
+                ],
+                "openai/gpt-4.1",
+            ),
+            # Test case 7: Cross-provider comparison - Anthropic is faster
+            (
+                "openai/gpt-5.2",
+                [
+                    "openai/gpt-5.2",  # Index 1 in OPENAI_MODEL_ORDER
+                    "anthropic/claude-haiku-4-5-20251001",  # Index 0 in ANTHROPIC_MODEL_ORDER
+                ],
+                "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}"

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

mito_ai/utils/model_utils.py

@@ -0,0 +1,178 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Tuple, Union, Optional, cast
+from mito_ai import constants
+from mito_ai.utils.version_utils import is_enterprise
+
+# Model ordering: [fastest, ..., slowest] for each provider
+ANTHROPIC_MODEL_ORDER = [
+    "claude-haiku-4-5-20251001",  # Fastest
+    "claude-sonnet-4-5-20250929",  # Slower
+]
+
+OPENAI_MODEL_ORDER = [
+    "gpt-4.1",      # Fastest
+    "gpt-5",
+    "gpt-5.2",      # Slower
+]
+
+GEMINI_MODEL_ORDER = [
+    "gemini-3-flash-preview",  # Fastest
+    "gemini-3-pro-preview",     # Slower
+]
+
+# Standard model names (used when not in enterprise mode or when LiteLLM is not configured)
+STANDARD_MODELS = [
+    "gpt-4.1",
+    "gpt-5.2",
+    "claude-sonnet-4-5-20250929",
+    "claude-haiku-4-5-20251001",
+    "gemini-3-flash-preview",
+    "gemini-3-pro-preview",
+]
+
+
+def get_available_models() -> List[str]:
+    """
+    Determine which models are available based on enterprise mode and LiteLLM configuration.
+    
+    Returns:
+        List of available model names. If enterprise mode is enabled AND LiteLLM is configured,
+        returns LiteLLM models. Otherwise, returns standard models.
+    """
+    # Check if enterprise mode is enabled AND LiteLLM is configured
+    if is_enterprise() and constants.LITELLM_BASE_URL and constants.LITELLM_MODELS:
+        # Return LiteLLM models (with provider prefixes)
+        return constants.LITELLM_MODELS
+    else:
+        # Return standard models
+        return STANDARD_MODELS
+
+
+def get_fast_model_for_selected_model(selected_model: str) -> str:
+    """
+    Get the fastest model for the client of the selected model.
+    
+    - For standard providers, returns the first (fastest) model from that provider's order.
+    - For LiteLLM models, finds the fastest available model from LiteLLM by comparing indices in the model order lists.
+    """
+    # Check if this is a LiteLLM model (has provider prefix like "openai/gpt-4o")
+    if "/" in selected_model:
+        
+        # Find the fastest model from available LiteLLM models
+        available_models = get_available_models()
+        if not available_models:
+            return selected_model
+        
+        # Filter to only LiteLLM models (those with "/") before splitting
+        litellm_models = [model for model in available_models if "/" in model]
+        if not litellm_models:
+            return selected_model
+        
+        available_provider_model_pairs: List[List[str]] = [model.split("/", 1) for model in litellm_models]
+        
+        # Get indices for all pairs and filter out None indices (unknown models)
+        pairs_with_indices = [(pair, get_model_order_index(pair)) for pair in available_provider_model_pairs]
+        valid_pairs_with_indices = [(pair, index) for pair, index in pairs_with_indices if index is not None]
+        
+        if not valid_pairs_with_indices:
+            return selected_model
+
+        # Find the pair with the minimum index (fastest model)
+        fastest_pair, _ = min(valid_pairs_with_indices, key=lambda x: x[1])
+        fastest_model = f"{fastest_pair[0]}/{fastest_pair[1]}"
+        
+        # If we found a fastest model, return it. Otherwise, use the selected model
+        if fastest_model:
+            return fastest_model
+        else:
+            return selected_model
+    
+    # Standard provider logic - ensure we return a model from the same provider
+    model_lower = selected_model.lower()
+    
+    # Determine provider and get fastest model
+    if model_lower.startswith('claude'):
+        return ANTHROPIC_MODEL_ORDER[0]
+    elif model_lower.startswith('gpt'):
+        return OPENAI_MODEL_ORDER[0]
+    elif model_lower.startswith('gemini'):
+        return GEMINI_MODEL_ORDER[0]
+
+    return selected_model
+
+def get_smartest_model_for_selected_model(selected_model: str) -> str:
+    """
+    Get the smartest model for the client of the selected model.
+    
+    - For standard providers, returns the last (smartest) model from that provider's order.
+    - For LiteLLM models, finds the smartest available model from LiteLLM by comparing indices in the model order lists.
+    """
+    # Check if this is a LiteLLM model (has provider prefix like "openai/gpt-4o")
+    if "/" in selected_model:
+        
+        # Extract provider from selected model
+        selected_provider, _ = selected_model.split("/", 1)
+        
+        # Find the smartest model from available LiteLLM models
+        available_models = get_available_models()
+        if not available_models:
+            return selected_model
+        
+        # Filter to only LiteLLM models (those with "/")
+        litellm_models = [model for model in available_models if "/" in model and model.startswith(f"{selected_provider}/")]
+        if not litellm_models:
+            return selected_model
+        
+        available_provider_model_pairs: List[List[str]] = [model.split("/", 1) for model in litellm_models]
+        
+        # Get indices for all pairs and filter out None indices (unknown models)
+        pairs_with_indices = [(pair, get_model_order_index(pair)) for pair in available_provider_model_pairs]
+        valid_pairs_with_indices = [(pair, index) for pair, index in pairs_with_indices if index is not None]
+        
+        if not valid_pairs_with_indices:
+            return selected_model
+
+        # Find the pair with the maximum index (smartest model)
+        smartest_pair, _ = max(valid_pairs_with_indices, key=lambda x: x[1])
+        smartest_model = f"{smartest_pair[0]}/{smartest_pair[1]}"
+        
+        # If we found a smartest model, return it. Otherwise, use the selected model
+        if smartest_model:
+            return smartest_model
+        else:
+            return selected_model
+    
+    # Standard provider logic
+    model_lower = selected_model.lower()
+    
+    # Determine provider and get smartest model
+    if model_lower.startswith('claude'):
+        return ANTHROPIC_MODEL_ORDER[-1]
+    elif model_lower.startswith('gpt'):
+        return OPENAI_MODEL_ORDER[-1]
+    elif model_lower.startswith('gemini'):
+        return GEMINI_MODEL_ORDER[-1]
+
+    return selected_model
+
+def get_model_order_index(pair: List[str]) -> Optional[int]:
+    provider, model_name = pair
+    if provider == "openai":
+        try:
+            return OPENAI_MODEL_ORDER.index(model_name)
+        except ValueError:
+            return None
+    elif provider == "anthropic":
+        try:
+            return ANTHROPIC_MODEL_ORDER.index(model_name)
+        except ValueError:
+            return None
+    elif provider == "google":
+        try:
+            return GEMINI_MODEL_ORDER.index(model_name)
+        except ValueError:
+            return None
+    else:
+        return None

mito_ai/utils/open_ai_utils.py

@@ -11,7 +11,6 @@ import json
 import time
 from typing import Any, Dict, List, Optional, Final, Union, AsyncGenerator, Tuple, Callable
 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 tornado.httpclient import AsyncHTTPClient
 from openai.types.chat import ChatCompletionMessageParam
 
@@ -153,19 +152,12 @@ async def stream_ai_completion_from_mito_server(
 
 
 def get_open_ai_completion_function_params(
-    message_type: MessageType,
     model: str, 
     messages: List[ChatCompletionMessageParam], 
     stream: bool,
     response_format_info: Optional[ResponseFormatInfo] = None,
 ) -> Dict[str, Any]:
     
-    print("MESSAGE TYPE: ", message_type)
-    message_requires_fast_model = does_message_require_fast_model(message_type)
-    model = FAST_OPENAI_MODEL if message_requires_fast_model else model
-    
-    print(f"model: {model}")
-    
     completion_function_params = {
         "model": model,
         "stream": stream,

mito_ai/utils/provider_utils.py

@@ -13,6 +13,12 @@ def get_model_provider(model: str) -> Union[str, None]:
     if not model:
         return None
 
+    # Check if model is a LiteLLM model (has provider prefix)
+    if "/" in model and any(
+        model.startswith(prefix) for prefix in ["openai/", "anthropic/", "google/", "ollama/"]
+    ):
+        return 'litellm'
+
     model_lower = model.lower()
 
     if model_lower.startswith('claude'):
@@ -25,32 +31,4 @@ def get_model_provider(model: str) -> Union[str, None]:
         return 'openai'
 
     return None
-
-
-def does_message_require_fast_model(message_type: MessageType) -> bool:
-    """
-    Determines if a message requires the fast model.
-    
-    The fast model is used for messages that are not chat messages.
-    For example, inline completions and chat name generation need to be fast
-    so they don't slow down the user's experience.
-    """
-    
-    if message_type in (
-        MessageType.CHAT,
-        MessageType.SMART_DEBUG,
-        MessageType.CODE_EXPLAIN,
-        MessageType.AGENT_EXECUTION,
-        MessageType.AGENT_SCRATCHPAD_RESULT,
-        MessageType.AGENT_AUTO_ERROR_FIXUP,
-    ):
-        return False
-    elif message_type in (MessageType.INLINE_COMPLETION, MessageType.CHAT_NAME_GENERATION):
-        return True
-    elif message_type in (MessageType.START_NEW_CHAT, MessageType.FETCH_HISTORY, MessageType.GET_THREADS, MessageType.DELETE_THREAD, MessageType.UPDATE_MODEL_CONFIG):
-        # These messages don't use any model, but we add them here for type safety
-        return True
-    else:
-        raise ValueError(f"Invalid message type: {message_type}")
-