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

mito_ai/tests/test_enterprise_mode.py

@@ -137,6 +137,61 @@ class TestModelValidation:
         
         from mito_ai.utils.model_utils import STANDARD_MODELS
         assert result == STANDARD_MODELS
+    
+    @patch('mito_ai.utils.model_utils.is_enterprise')
+    @patch('mito_ai.utils.model_utils.constants')
+    @patch('mito_ai.utils.model_utils.is_abacus_configured')
+    def test_provider_manager_validates_abacus_model(self, mock_is_abacus_configured, mock_constants, mock_is_enterprise, provider_config: Config):
+        """Test that ProviderManager validates Abacus models against available models."""
+        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/gpt-5.2"]
+        
+        provider_manager = ProviderManager(config=provider_config)
+        provider_manager.set_selected_model("Abacus/gpt-4.1")
+        
+        # Should not raise an error for valid model
+        available_models = get_available_models()
+        assert "Abacus/gpt-4.1" in available_models
+    
+    @patch('mito_ai.utils.model_utils.is_enterprise')
+    @patch('mito_ai.utils.model_utils.constants')
+    @patch('mito_ai.utils.model_utils.is_abacus_configured')
+    @pytest.mark.asyncio
+    async def test_provider_manager_rejects_invalid_abacus_model(self, mock_is_abacus_configured, mock_constants, mock_is_enterprise, provider_config: Config):
+        """Test that ProviderManager rejects invalid Abacus models."""
+        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"]
+        mock_constants.ABACUS_API_KEY = "test-key"
+        
+        provider_manager = ProviderManager(config=provider_config)
+        provider_manager.set_selected_model("invalid-model")
+        
+        messages: list[ChatCompletionMessageParam] = [{"role": "user", "content": "test"}]
+        
+        # Should raise ValueError for invalid model
+        with pytest.raises(ValueError, match="is not in the allowed model list"):
+            await provider_manager.request_completions(
+                message_type=MessageType.CHAT,
+                messages=messages
+            )
+    
+    @patch('mito_ai.utils.model_utils.is_enterprise')
+    @patch('mito_ai.utils.model_utils.constants')
+    @patch('mito_ai.utils.model_utils.is_abacus_configured')
+    def test_available_models_endpoint_returns_abacus_models(self, mock_is_abacus_configured, mock_constants, mock_is_enterprise):
+        """Test that /available-models endpoint returns Abacus models when configured."""
+        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 TestModelStorage:

mito_ai/tests/test_model_utils.py

@@ -22,11 +22,11 @@ class TestGetAvailableModels:
         """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"]
+        mock_constants.LITELLM_MODELS = ["litellm/openai/gpt-4o", "litellm/anthropic/claude-3-5-sonnet"]
         
         result = get_available_models()
         
-        assert result == ["openai/gpt-4o", "anthropic/claude-3-5-sonnet"]
+        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')
@@ -56,7 +56,7 @@ class TestGetAvailableModels:
         """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"]
+        mock_constants.LITELLM_MODELS = ["litellm/openai/gpt-4o"]
         
         result = get_available_models()
         
@@ -73,6 +73,20 @@ class TestGetAvailableModels:
         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:
@@ -114,51 +128,51 @@ class TestGetFastModelForSelectedModel:
         [
             # 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",
+                "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
             (
-                "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",
+                "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
             (
-                "google/gemini-3-pro-preview",
-                ["google/gemini-3-pro-preview", "google/gemini-3-flash-preview"],
-                "google/gemini-3-flash-preview",
+                "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",
-                ["openai/gpt-4.1", "unknown/provider/model"],
-                "openai/gpt-4.1",
+                ["litellm/openai/gpt-4.1", "unknown/provider/model"],
+                "litellm/openai/gpt-4.1",
             ),
             # Test case 5: Single LiteLLM model returns itself
             (
-                "openai/gpt-4o",
-                ["openai/gpt-4o"],
-                "openai/gpt-4o",
+                "litellm/openai/gpt-4o",
+                ["litellm/openai/gpt-4o"],
+                "litellm/openai/gpt-4o",
             ),
             # Test case 6: Cross-provider comparison - OpenAI is faster
             (
-                "anthropic/claude-sonnet-4-5-20250929",
+                "litellm/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
+                    "litellm/openai/gpt-4.1",  # Index 0 in OPENAI_MODEL_ORDER
+                    "litellm/anthropic/claude-sonnet-4-5-20250929",  # Index 1 in ANTHROPIC_MODEL_ORDER
                 ],
-                "openai/gpt-4.1",
+                "litellm/openai/gpt-4.1",
             ),
             # Test case 7: Cross-provider comparison - Anthropic is faster
             (
-                "openai/gpt-5.2",
+                "litellm/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
+                    "litellm/openai/gpt-5.2",  # Index 1 in OPENAI_MODEL_ORDER
+                    "litellm/anthropic/claude-haiku-4-5-20251001",  # Index 0 in ANTHROPIC_MODEL_ORDER
                 ],
-                "anthropic/claude-haiku-4-5-20251001",
+                "litellm/anthropic/claude-haiku-4-5-20251001",
             ),
         ],
         ids=[
@@ -269,3 +283,80 @@ class TestGetFastModelForSelectedModel:
         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/model_utils.py

@@ -4,6 +4,7 @@
 from typing import List, Tuple, Union, Optional, cast
 from mito_ai import constants
 from mito_ai.utils.version_utils import is_enterprise
+from mito_ai.enterprise.utils import is_abacus_configured
 
 # Model ordering: [fastest, ..., slowest] for each provider
 ANTHROPIC_MODEL_ORDER = [
@@ -35,15 +36,23 @@ STANDARD_MODELS = [
 
 def get_available_models() -> List[str]:
     """
-    Determine which models are available based on enterprise mode and LiteLLM configuration.
+    Determine which models are available based on enterprise mode and router configuration.
+    
+    Priority order:
+    1. Abacus (if configured)
+    2. LiteLLM (if configured)
+    3. Standard models
     
     Returns:
-        List of available model names. If enterprise mode is enabled AND LiteLLM is configured,
-        returns LiteLLM models. Otherwise, returns standard models.
+        List of available model names with appropriate prefixes.
     """
+    # Check if enterprise mode is enabled AND Abacus is configured (highest priority)
+    if is_abacus_configured():
+        # Return Abacus models (with Abacus/ prefix)
+        return constants.ABACUS_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)
+    elif is_enterprise() and constants.LITELLM_BASE_URL and constants.LITELLM_MODELS:
+        # Return LiteLLM models (with LiteLLM/provider/ prefix or legacy provider/ prefix)
         return constants.LITELLM_MODELS
     else:
         # Return standard models
@@ -55,39 +64,50 @@ 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.
+    - For enterprise router models (Abacus/LiteLLM), finds the fastest available model by comparing indices.
     """
-    # 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
+    # Check if this is an enterprise router model (has "/" or router prefix)
+    if "/" in selected_model or selected_model.lower().startswith(('abacus/', 'litellm/')):
+        # Find the fastest model from available 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:
+        # Filter to only router models (those with "/")
+        router_models = [model for model in available_models if "/" in model]
+        if not router_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]
+        # Extract provider/model pairs for ordering
+        pairs_with_indices = []
+        for model in router_models:
+            # Strip router prefix to get underlying model info
+            model_without_router = strip_router_prefix(model)
+            
+            # For Abacus: model_without_router is just the model name (e.g., "gpt-4.1")
+            # For LiteLLM: model_without_router is "provider/model" (e.g., "openai/gpt-4.1")
+            if "/" in model_without_router:
+                # LiteLLM format: provider/model
+                pair = model_without_router.split("/", 1)
+            else:
+                # Abacus format: just model name, need to determine provider
+                provider = get_underlying_model_provider(model)
+                if provider:
+                    pair = [provider, model_without_router]
+                else:
+                    continue
+            
+            index = get_model_order_index(pair)
+            if index is not None:
+                pairs_with_indices.append((model, index))
         
-        if not valid_pairs_with_indices:
+        if not 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]}"
+        # Find the model with the minimum index (fastest model)
+        fastest_model, _ = min(pairs_with_indices, key=lambda x: x[1])
         
-        # If we found a fastest model, return it. Otherwise, use the selected model
-        if fastest_model:
-            return fastest_model
-        else:
-            return selected_model
+        return fastest_model
     
     # Standard provider logic - ensure we return a model from the same provider
     model_lower = selected_model.lower()
@@ -107,42 +127,57 @@ 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.
+    - For enterprise router models (Abacus/LiteLLM), finds the smartest available model by comparing indices.
     """
-    # 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)
+    # Check if this is an enterprise router model (has "/" or router prefix)
+    if "/" in selected_model or selected_model.lower().startswith(('abacus/', 'litellm/')):
+        # Extract underlying provider from selected model
+        selected_provider = get_underlying_model_provider(selected_model)
+        if not selected_provider:
+            return selected_model
         
-        # Find the smartest model from available LiteLLM models
+        # Find the smartest model from available 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
+        # Filter to only router models with the same underlying provider
+        router_models = []
+        for model in available_models:
+            if "/" in model:
+                model_provider = get_underlying_model_provider(model)
+                if model_provider == selected_provider:
+                    router_models.append(model)
         
-        available_provider_model_pairs: List[List[str]] = [model.split("/", 1) for model in litellm_models]
+        if not router_models:
+            return selected_model
         
-        # 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]
+        # Extract provider/model pairs for ordering
+        pairs_with_indices = []
+        for model in router_models:
+            # Strip router prefix to get underlying model info
+            model_without_router = strip_router_prefix(model)
+            
+            # For Abacus: model_without_router is just the model name (e.g., "gpt-4.1")
+            # For LiteLLM: model_without_router is "provider/model" (e.g., "openai/gpt-4.1")
+            if "/" in model_without_router:
+                # LiteLLM format: provider/model
+                pair = model_without_router.split("/", 1)
+            else:
+                # Abacus format: just model name, provider already determined
+                pair = [selected_provider, model_without_router]
+            
+            index = get_model_order_index(pair)
+            if index is not None:
+                pairs_with_indices.append((model, index))
         
-        if not valid_pairs_with_indices:
+        if not 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]}"
+        # Find the model with the maximum index (smartest model)
+        smartest_model, _ = max(pairs_with_indices, key=lambda x: x[1])
         
-        # If we found a smartest model, return it. Otherwise, use the selected model
-        if smartest_model:
-            return smartest_model
-        else:
-            return selected_model
+        return smartest_model
     
     # Standard provider logic
     model_lower = selected_model.lower()
@@ -157,6 +192,50 @@ def get_smartest_model_for_selected_model(selected_model: str) -> str:
 
     return selected_model
 
+def strip_router_prefix(model: str) -> str:
+    """
+    Strip router prefix from model name.
+    
+    Examples:
+    - "Abacus/gpt-4.1" -> "gpt-4.1"
+    - "LiteLLM/openai/gpt-4.1" -> "openai/gpt-4.1"
+    - "gpt-4.1" -> "gpt-4.1" (no prefix, return as-is)
+    """
+    if model.lower().startswith('abacus/'):
+        return model[7:]  # Strip "Abacus/"
+    elif model.lower().startswith('litellm/'):
+        return model[8:]  # Strip "LiteLLM/"
+    return model
+
+def get_underlying_model_provider(full_model_provider_id: str) -> Optional[str]:
+    """
+    Determine the underlying AI provider from a model identifier.
+    
+    For Abacus models (Abacus/model), determine the provider from model name pattern.
+    For LiteLLM models (LiteLLM/provider/model), extract the provider from the prefix.
+    
+    Returns:
+        Provider name ("openai", "anthropic", "google") or None if cannot determine.
+    """
+    # Strip router prefix first
+    model_without_router = strip_router_prefix(full_model_provider_id)
+    
+    # Check if it's a LiteLLM format (provider/model)
+    if "/" in model_without_router:
+        provider, _ = model_without_router.split("/", 1)
+        return provider.lower()
+    
+    # For Abacus models without provider prefix, determine from model name
+    model_lower = model_without_router.lower()
+    if model_lower.startswith('gpt'):
+        return 'openai'
+    elif model_lower.startswith('claude'):
+        return 'anthropic'
+    elif model_lower.startswith('gemini'):
+        return 'google'
+    
+    return None
+
 def get_model_order_index(pair: List[str]) -> Optional[int]:
     provider, model_name = pair
     if provider == "openai":

mito_ai/utils/open_ai_utils.py

@@ -6,30 +6,18 @@
 
 # Copyright (c) Saga Inc.
 
-import asyncio
-import json
-import time
-from typing import Any, Dict, List, Optional, Final, Union, AsyncGenerator, Tuple, Callable
+from typing import Any, Dict, List, Optional, Union, AsyncGenerator, Tuple, Callable
 from mito_ai.utils.mito_server_utils import get_response_from_mito_server, stream_response_from_mito_server
-from tornado.httpclient import AsyncHTTPClient
 from openai.types.chat import ChatCompletionMessageParam
-
-from mito_ai.utils.utils import is_running_test
-from mito_ai.completions.models import MessageType, ResponseFormatInfo, CompletionReply, CompletionStreamChunk, CompletionItem
+from mito_ai.completions.models import 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
-from mito_ai.utils.version_utils import is_pro
-from mito_ai.utils.server_limits import check_mito_server_quota
-from mito_ai.utils.telemetry_utils import log_ai_completion_success
-from .utils import _create_http_client
+from mito_ai.enterprise.utils import is_abacus_configured
 from mito_ai.constants import MITO_OPENAI_URL
 
-
 __user_email: Optional[str] = None
 __user_id: Optional[str] = None
 
-FAST_OPENAI_MODEL = "gpt-4.1-nano"
-
 def _prepare_request_data_and_headers(
     last_message_content: Union[str, None],
     ai_completion_data: Dict[str, Any],
@@ -171,25 +159,33 @@ def get_open_ai_completion_function_params(
     # Pydantic models are supported by the OpenAI API, however, we need to be able to 
     # serialize it for requests that are going to be sent to the mito server. 
     # OpenAI expects a very specific schema as seen below. 
+    # Note: Abacus only supports {"type": "json"} format, not the full JSON schema format.
     if response_format_info:
-        json_schema = response_format_info.format.schema()
-        
-        # Add additionalProperties: False to the top-level schema
-        json_schema["additionalProperties"] = False
-        
-        # Nested object definitions in $defs need to have additionalProperties set to False also
-        if "$defs" in json_schema:
-            for def_name, def_schema in json_schema["$defs"].items():
-                if def_schema.get("type") == "object":
-                    def_schema["additionalProperties"] = False
-        
-        completion_function_params["response_format"] = {
-            "type": "json_schema",
-            "json_schema": {
-                "name": f"{response_format_info.name}",
-                "schema": json_schema,
-                "strict": True
+        # Check if we're using Abacus - it only supports simple {"type": "json"} format
+        if is_abacus_configured() or model.lower().startswith('abacus/'):
+            completion_function_params["response_format"] = {
+                "type": "json"
+            }
+        else:
+            # For OpenAI and other providers, use the full JSON schema format
+            json_schema = response_format_info.format.schema()
+            
+            # Add additionalProperties: False to the top-level schema
+            json_schema["additionalProperties"] = False
+            
+            # Nested object definitions in $defs need to have additionalProperties set to False also
+            if "$defs" in json_schema:
+                for def_name, def_schema in json_schema["$defs"].items():
+                    if def_schema.get("type") == "object":
+                        def_schema["additionalProperties"] = False
+            
+            completion_function_params["response_format"] = {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": f"{response_format_info.name}",
+                    "schema": json_schema,
+                    "strict": True
+                }
             }
-        }
 
     return completion_function_params

mito_ai/utils/provider_utils.py

@@ -8,19 +8,25 @@ from mito_ai.completions.models import MessageType
 
 def get_model_provider(model: str) -> Union[str, None]:
     """
-    Determine the model type based on the model name prefix
+    Determine the model type based on the model name prefix.
+    
+    Priority order:
+    1. Check for router prefixes (Abacus/, LiteLLM/)
+    2. Check for legacy LiteLLM format (provider/model)
+    3. Check for standard model name patterns
     """
     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()
 
+    # Check for router prefixes first (highest priority)
+    if model_lower.startswith('abacus/'):
+        return 'abacus'
+    elif model_lower.startswith('litellm/'):
+        return 'litellm'    
+
+    # Check for standard model name patterns
     if model_lower.startswith('claude'):
         return 'claude'
     elif model_lower.startswith('gemini'):

{mito_ai-0.1.58.data → mito_ai-0.1.59.data}/data/share/jupyter/labextensions/mito_ai/build_log.json

@@ -720,7 +720,7 @@
             "semver": {},
             "vscode-diff": {},
             "mito_ai": {
-              "version": "0.1.58",
+              "version": "0.1.59",
               "singleton": true,
               "import": "/home/runner/work/mito/mito/mito-ai/lib/index.js"
             }

{mito_ai-0.1.58.data → mito_ai-0.1.59.data}/data/share/jupyter/labextensions/mito_ai/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.58",
+  "version": "0.1.59",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",
@@ -141,7 +141,7 @@
     "schemaDir": "schema",
     "themePath": "style/theme/theme.css",
     "_build": {
-      "load": "static/remoteEntry.570df809a692f53a7ab7.js",
+      "load": "static/remoteEntry.f7decebaf69618541e0f.js",
       "extension": "./extension",
       "style": "./style"
     }