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

mito_ai/__init__.py

@@ -9,6 +9,7 @@ from mito_ai.completions.message_history import GlobalMessageHistory
 from mito_ai.app_deploy.handlers import AppDeployHandler
 from mito_ai.log.urls import get_log_urls
 from mito_ai.utils.litellm_utils import is_litellm_configured
+from mito_ai.enterprise.utils import is_abacus_configured
 from mito_ai.version_check import VersionCheckHandler
 from mito_ai.db.urls import get_db_urls
 from mito_ai.settings.urls import get_settings_urls
@@ -101,10 +102,12 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
 
     web_app.add_handlers(host_pattern, handlers)
     
-    # Log enterprise mode status and LiteLLM configuration
+    # Log enterprise mode status and router configuration
     if is_enterprise():
         server_app.log.info("Enterprise mode enabled")
-        if is_litellm_configured():
+        if is_abacus_configured():
+            server_app.log.info(f"Abacus AI configured: endpoint={constants.ABACUS_BASE_URL}, models={constants.ABACUS_MODELS}")
+        elif is_litellm_configured():
             server_app.log.info(f"LiteLLM configured: endpoint={constants.LITELLM_BASE_URL}, models={constants.LITELLM_MODELS}")
     
     server_app.log.info("Loaded the mito_ai server extension")

mito_ai/_version.py

@@ -1,4 +1,4 @@
 # This file is auto-generated by Hatchling. As such, do not:
 #   - modify
 #   - track in version control e.g. be sure to add to .gitignore
-__version__ = VERSION = '0.1.58'
+__version__ = VERSION = '0.1.59'

mito_ai/completions/prompt_builders/prompt_constants.py

@@ -20,6 +20,8 @@ Rules:
 - All imports must appear at the top, before the chart configuration section.
 - Variables with multiple words should be underscore-separated.
 - All colors should be in hex format (e.g., "#3498db"). Use quotes around the hex string: COLOR = "#3498db" or COLOR = '#3498db'. Do NOT nest quotes.
+- Never use RGB/RGBA tuples/lists for colors (e.g. (0, 0.4, 0.8, 0.8) is forbidden).
+- If transparency is needed, store it separately as ALPHA = 0.8 and apply it in code (e.g. to_rgba(HEX_COLOR, ALPHA)).
 - Variables can only be strings, numbers, booleans, tuples, or lists. 
 - NEVER include comments on the same line as a variable assignment. Each variable assignment must be on its own line with no trailing comments.
 - For string values, use either single or double quotes (e.g., TITLE = "Sales by Product" or TITLE = 'Sales by Product'). Do not use nested quotes (e.g., do NOT use '"value"').

mito_ai/constants.py

@@ -2,7 +2,7 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 import os
-from typing import Union
+from typing import Union, List
 
 # Claude
 ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY")
@@ -23,12 +23,34 @@ AZURE_OPENAI_API_VERSION = os.environ.get("AZURE_OPENAI_API_VERSION")
 AZURE_OPENAI_ENDPOINT = os.environ.get("AZURE_OPENAI_ENDPOINT")
 AZURE_OPENAI_MODEL = os.environ.get("AZURE_OPENAI_MODEL")
 
+def parse_comma_separated_models(models_str: str) -> List[str]:
+    """
+    Parse a comma-separated string of model names into a list.
+    Handles quoted and unquoted values, stripping whitespace and quotes.
+    
+    Args:
+        models_str: Comma-separated string of model names (e.g., "model1,model2" or '"model1","model2"')
+    
+    Returns:
+        List of model names with whitespace and quotes stripped
+    """
+    if not models_str:
+        return []
+    return [model.strip().strip('"\'') for model in models_str.split(",") if model.strip()]
+
 # LiteLLM Config (Enterprise mode only)
 LITELLM_BASE_URL = os.environ.get("LITELLM_BASE_URL")
 LITELLM_API_KEY = os.environ.get("LITELLM_API_KEY")
 LITELLM_MODELS_STR = os.environ.get("LITELLM_MODELS", "")
-# Parse comma-separated string into list, strip whitespace
-LITELLM_MODELS = [model.strip() for model in LITELLM_MODELS_STR.split(",") if model.strip()] if LITELLM_MODELS_STR else []
+# Parse comma-separated string into list, strip whitespace and quotes
+LITELLM_MODELS = parse_comma_separated_models(LITELLM_MODELS_STR)
+
+# Abacus AI Config (Enterprise mode only)
+ABACUS_BASE_URL = os.environ.get("ABACUS_BASE_URL")
+ABACUS_API_KEY = os.environ.get("ABACUS_API_KEY")
+ABACUS_MODELS_STR = os.environ.get("ABACUS_MODELS", "")
+# Parse comma-separated string into list, strip whitespace and quotes
+ABACUS_MODELS = parse_comma_separated_models(ABACUS_MODELS_STR)
 
 # Mito AI Base URLs and Endpoint Paths
 MITO_PROD_BASE_URL = "https://7eax4i53f5odkshhlry4gw23by0yvnuv.lambda-url.us-east-1.on.aws/v2"

mito_ai/enterprise/litellm_client.py

@@ -11,6 +11,7 @@ from mito_ai.completions.models import (
     CompletionItem,
 )
 from mito_ai.utils.litellm_utils import get_litellm_completion_function_params
+from mito_ai.utils.model_utils import strip_router_prefix
 import litellm
 
 class LiteLLMClient:
@@ -28,7 +29,7 @@ class LiteLLMClient:
     async def request_completions(
         self,
         messages: List[ChatCompletionMessageParam],
-        model: str,  # Should include provider prefix (e.g., "openai/gpt-4o")
+        model: str,  # Should include provider prefix (e.g., "LiteLLM/openai/gpt-4o")
         response_format_info: Optional[ResponseFormatInfo] = None,
         message_type: MessageType = MessageType.CHAT
     ) -> str:
@@ -37,16 +38,19 @@ class LiteLLMClient:
         
         Args:
             messages: List of chat messages
-            model: Model name with provider prefix (e.g., "openai/gpt-4o")
+            model: Model name with router and provider prefix (e.g., "LiteLLM/openai/gpt-4o")
             response_format_info: Optional response format specification
             message_type: Type of message (chat, agent execution, etc.)
             
         Returns:
             The completion text response
         """
+        # Strip router prefix if present (LiteLLM/ prefix)
+        model_for_litellm = strip_router_prefix(model)
+        
         # Prepare parameters for LiteLLM
         params = get_litellm_completion_function_params(
-            model=model,
+            model=model_for_litellm,
             messages=messages,
             api_key=self.api_key,
             api_base=self.base_url,
@@ -82,7 +86,7 @@ class LiteLLMClient:
         
         Args:
             messages: List of chat messages
-            model: Model name with provider prefix (e.g., "openai/gpt-4o")
+            model: Model name with router and provider prefix (e.g., "LiteLLM/openai/gpt-4o")
             message_type: Type of message (chat, agent execution, etc.)
             message_id: ID of the message being processed
             reply_fn: Function to call with each chunk for streaming replies
@@ -93,9 +97,12 @@ class LiteLLMClient:
         """
         accumulated_response = ""
         
+        # Strip router prefix if present (LiteLLM/ prefix)
+        model_for_litellm = strip_router_prefix(model)
+        
         # Prepare parameters for LiteLLM
         params = get_litellm_completion_function_params(
-            model=model,
+            model=model_for_litellm,
             messages=messages,
             api_key=self.api_key,
             api_base=self.base_url,

mito_ai/enterprise/utils.py

@@ -5,11 +5,25 @@
 # Distributed under the terms of the The Mito Enterprise license.
 
 from mito_ai.utils.version_utils import is_enterprise, is_mitosheet_private
-from mito_ai.constants import AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_API_VERSION, AZURE_OPENAI_MODEL
+from mito_ai.constants import (
+    AZURE_OPENAI_API_KEY, 
+    AZURE_OPENAI_ENDPOINT, 
+    AZURE_OPENAI_API_VERSION, 
+    AZURE_OPENAI_MODEL,
+    ABACUS_BASE_URL,
+    ABACUS_MODELS
+)
 
 def is_azure_openai_configured() -> bool:
     """
     Azure OpenAI is only supported for Mito Enterprise users
     """
     is_allowed_to_use_azure = is_enterprise() or is_mitosheet_private()
-    return all([is_allowed_to_use_azure, AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_API_VERSION, AZURE_OPENAI_MODEL])
+    return all([is_allowed_to_use_azure, AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_API_VERSION, AZURE_OPENAI_MODEL])
+
+def is_abacus_configured() -> bool:
+    """
+    Abacus AI is only supported for Mito Enterprise users.
+    Checks if Abacus AI is configured with base URL and models.
+    """
+    return all([is_enterprise(), ABACUS_BASE_URL, ABACUS_MODELS])

mito_ai/openai_client.py

@@ -11,8 +11,9 @@ from traitlets import Instance, default, validate
 from traitlets.config import LoggingConfigurable
 
 from mito_ai import constants
-from mito_ai.enterprise.utils import is_azure_openai_configured
+from mito_ai.enterprise.utils import is_azure_openai_configured, is_abacus_configured
 from mito_ai.logger import get_logger
+from mito_ai.utils.model_utils import strip_router_prefix
 from mito_ai.completions.models import (
     AICapabilities,
     CompletionError,
@@ -24,12 +25,11 @@ from mito_ai.completions.models import (
     ResponseFormatInfo,
 )
 from mito_ai.utils.open_ai_utils import (
-    check_mito_server_quota,
     get_ai_completion_from_mito_server,
     get_open_ai_completion_function_params,
     stream_ai_completion_from_mito_server,
 )
-from mito_ai.utils.server_limits import update_mito_server_quota
+from mito_ai.utils.server_limits import update_mito_server_quota, check_mito_server_quota
 
 OPENAI_MODEL_FALLBACK = "gpt-4.1"
 
@@ -68,6 +68,14 @@ This attribute is observed by the websocket provider to push the error to the cl
                 provider="Azure OpenAI",
             )
 
+        if is_abacus_configured():
+            return AICapabilities(
+                configuration={
+                    "model": "<dynamic>"
+                },
+                provider="Abacus AI",
+            )
+
         if constants.OLLAMA_MODEL:
             return AICapabilities(
                 configuration={
@@ -121,6 +129,10 @@ This attribute is observed by the websocket provider to push the error to the cl
                 timeout=self.timeout,
             )
         
+        elif is_abacus_configured():
+            base_url = constants.ABACUS_BASE_URL
+            llm_api_key = constants.ABACUS_API_KEY
+            self.log.debug(f"Using Abacus AI with base URL: {constants.ABACUS_BASE_URL}")
         elif constants.OLLAMA_MODEL:
             base_url = constants.OLLAMA_BASE_URL
             llm_api_key = "ollama"
@@ -141,17 +153,25 @@ This attribute is observed by the websocket provider to push the error to the cl
         )
         return client
 
-    def _adjust_model_for_azure_or_ollama(self, model: str) -> str:
+    def _adjust_model_for_provider(self, model: str) -> str:
         
         # If they have set an Azure OpenAI model, then we always use it
         if is_azure_openai_configured() and constants.AZURE_OPENAI_MODEL is not None:
             self.log.debug(f"Resolving to Azure OpenAI model: {constants.AZURE_OPENAI_MODEL}")
+            # TODO: We should update Azure so it works the way LiteLLM and Abacus do: 
+            # when configured, we only show models from Azure in the UI. 
             return constants.AZURE_OPENAI_MODEL
         
         # If they have set an Ollama model, then we use it
         if constants.OLLAMA_MODEL is not None:
             return constants.OLLAMA_MODEL
         
+        # If using Abacus, strip the "Abacus/" prefix from the model name
+        if is_abacus_configured() and model.lower().startswith('abacus/'):
+            stripped_model = strip_router_prefix(model)
+            self.log.debug(f"Stripping Abacus prefix: {model} -> {stripped_model}")
+            return stripped_model
+        
         # Otherwise, we use the model they provided
         return model
         
@@ -186,7 +206,7 @@ This attribute is observed by the websocket provider to push the error to the cl
         )
         
         # If they have set an Azure OpenAI or Ollama model, then we use it
-        completion_function_params["model"] = self._adjust_model_for_azure_or_ollama(completion_function_params["model"])
+        completion_function_params["model"] = self._adjust_model_for_provider(completion_function_params["model"])
 
         if self._active_async_client is not None:
             response = await self._active_async_client.chat.completions.create(**completion_function_params)
@@ -236,7 +256,7 @@ This attribute is observed by the websocket provider to push the error to the cl
             model, messages, True, response_format_info
         )
         
-        completion_function_params["model"] = self._adjust_model_for_azure_or_ollama(completion_function_params["model"])
+        completion_function_params["model"] = self._adjust_model_for_provider(completion_function_params["model"])
 
         try:
             if self._active_async_client is not None:

mito_ai/provider_manager.py

@@ -27,6 +27,7 @@ from mito_ai.completions.models import (
     ResponseFormatInfo,
 )
 from mito_ai.utils.litellm_utils import is_litellm_configured
+from mito_ai.enterprise.utils import is_abacus_configured
 from mito_ai.utils.telemetry_utils import (
     MITO_SERVER_KEY,
     USER_KEY,
@@ -79,6 +80,12 @@ This attribute is observed by the websocket provider to push the error to the cl
         # TODO: We should validate that these keys are actually valid for the provider
         # otherwise it will look like we are using the user_key when actually falling back 
         # to the mito server because the key is invalid. 
+        if is_abacus_configured():
+            return AICapabilities(
+                configuration={"model": "<dynamic>"},
+                provider="Abacus AI",
+            )
+            
         if is_litellm_configured():
             return AICapabilities(
                 configuration={"model": "<dynamic>"},
@@ -116,6 +123,9 @@ This attribute is observed by the websocket provider to push the error to the cl
         # TODO: We should validate that these keys are actually valid for the provider
         # otherwise it will look like we are using the user_key when actually falling back 
         # to the mito server because the key is invalid. 
+        if is_abacus_configured():
+            return USER_KEY
+            
         if is_litellm_configured():
             return USER_KEY
         
@@ -172,7 +182,16 @@ This attribute is observed by the websocket provider to push the error to the cl
         # Retry loop
         for attempt in range(max_retries + 1):
             try:
-                if model_type == "litellm":
+                if model_type == "abacus":
+                    if not self._openai_client:
+                        raise RuntimeError("OpenAI client is not initialized.")
+                    completion = await self._openai_client.request_completions(
+                        message_type=message_type,
+                        messages=messages,
+                        model=resolved_model,
+                        response_format_info=response_format_info
+                    )
+                elif model_type == "litellm":
                     from mito_ai.enterprise.litellm_client import LiteLLMClient
                     if not constants.LITELLM_BASE_URL:
                         raise ValueError("LITELLM_BASE_URL is required for LiteLLM models")
@@ -299,7 +318,20 @@ This attribute is observed by the websocket provider to push the error to the cl
         ))
 
         try:
-            if model_type == "litellm":
+            if model_type == "abacus":
+                if not self._openai_client:
+                    raise RuntimeError("OpenAI client is not initialized.")
+                accumulated_response = await self._openai_client.stream_completions(
+                    message_type=message_type,
+                    messages=messages,
+                    model=resolved_model,
+                    message_id=message_id,
+                    thread_id=thread_id,
+                    reply_fn=reply_fn,
+                    user_input=user_input,
+                    response_format_info=response_format_info
+                )
+            elif model_type == "litellm":
                 from mito_ai.enterprise.litellm_client import LiteLLMClient
                 if not constants.LITELLM_BASE_URL:
                     raise ValueError("LITELLM_BASE_URL is required for LiteLLM models")

mito_ai/tests/message_history/test_generate_short_chat_name.py

@@ -23,7 +23,8 @@ PROVIDER_TEST_CASES = [
     ("gpt-4.1", "mito_ai.provider_manager.OpenAIClient"),
     ("claude-sonnet-4-5-20250929", "mito_ai.provider_manager.AnthropicClient"),
     ("gemini-3-flash-preview", "mito_ai.provider_manager.GeminiClient"),
-    ("openai/gpt-4o", "mito_ai.provider_manager.LiteLLMClient"),  # LiteLLM test case
+    ("litellm/openai/gpt-4o", "mito_ai.provider_manager.LiteLLMClient"),  # LiteLLM test case
+    ("Abacus/gpt-4.1", "mito_ai.provider_manager.OpenAIClient"),  # Abacus test case (uses OpenAIClient)
 ]
 
 @pytest.mark.parametrize("selected_model,client_patch_path", PROVIDER_TEST_CASES)
@@ -49,13 +50,27 @@ async def test_generate_short_chat_name_uses_correct_provider_and_fast_model(
         # Patch constants both at the source and where they're imported in model_utils
         monkeypatch.setattr("mito_ai.constants.LITELLM_BASE_URL", "https://litellm-server.com")
         monkeypatch.setattr("mito_ai.constants.LITELLM_API_KEY", "fake-litellm-key")
-        monkeypatch.setattr("mito_ai.constants.LITELLM_MODELS", ["openai/gpt-4o", "anthropic/claude-3-5-sonnet"])
+        monkeypatch.setattr("mito_ai.constants.LITELLM_MODELS", ["litellm/openai/gpt-4o", "litellm/anthropic/claude-3-5-sonnet"])
         # Also patch where constants is imported in model_utils (where get_available_models uses it)
         monkeypatch.setattr("mito_ai.utils.model_utils.constants.LITELLM_BASE_URL", "https://litellm-server.com")
-        monkeypatch.setattr("mito_ai.utils.model_utils.constants.LITELLM_MODELS", ["openai/gpt-4o", "anthropic/claude-3-5-sonnet"])
+        monkeypatch.setattr("mito_ai.utils.model_utils.constants.LITELLM_MODELS", ["litellm/openai/gpt-4o", "litellm/anthropic/claude-3-5-sonnet"])
         # Mock is_enterprise to return True so LiteLLM models are available
         monkeypatch.setattr("mito_ai.utils.version_utils.is_enterprise", lambda: True)
     
+    # Set up Abacus constants if testing Abacus
+    if selected_model.startswith("Abacus/"):
+        # Patch constants both at the source and where they're imported in model_utils
+        monkeypatch.setattr("mito_ai.constants.ABACUS_BASE_URL", "https://routellm.abacus.ai/v1")
+        monkeypatch.setattr("mito_ai.constants.ABACUS_API_KEY", "fake-abacus-key")
+        monkeypatch.setattr("mito_ai.constants.ABACUS_MODELS", ["Abacus/gpt-4.1", "Abacus/claude-haiku-4-5-20251001"])
+        # Also patch where constants is imported in model_utils (where get_available_models uses it)
+        monkeypatch.setattr("mito_ai.utils.model_utils.constants.ABACUS_BASE_URL", "https://routellm.abacus.ai/v1")
+        monkeypatch.setattr("mito_ai.utils.model_utils.constants.ABACUS_MODELS", ["Abacus/gpt-4.1", "Abacus/claude-haiku-4-5-20251001"])
+        # Mock is_abacus_configured to return True so Abacus models are available
+        monkeypatch.setattr("mito_ai.utils.model_utils.is_abacus_configured", lambda: True)
+        # Mock is_enterprise to return True so enterprise models are available
+        monkeypatch.setattr("mito_ai.utils.version_utils.is_enterprise", lambda: True)
+    
     # Create mock client for the specific provider being tested
     mock_client = MagicMock()
     mock_client.request_completions = AsyncMock(return_value="Test Chat Name")
@@ -87,12 +102,28 @@ async def test_generate_short_chat_name_uses_correct_provider_and_fast_model(
         # Patch LiteLLMClient where it's defined (it's imported inside request_completions)
         # Also patch get_available_models to return LiteLLM models
         with patch("mito_ai.enterprise.litellm_client.LiteLLMClient", return_value=mock_client), \
-             patch("mito_ai.provider_manager.get_available_models", return_value=["openai/gpt-4o", "anthropic/claude-3-5-sonnet"]):
+             patch("mito_ai.provider_manager.get_available_models", return_value=["litellm/openai/gpt-4o", "litellm/anthropic/claude-3-5-sonnet"]):
+            result = await generate_short_chat_name(
+                user_message="What is the capital of France?",
+                assistant_message="The capital of France is Paris.",
+                llm_provider=llm_provider
+            )
+    elif selected_model.startswith("Abacus/"):
+        # For Abacus, it uses OpenAIClient, so patch the instance's method
+        # Also patch get_available_models to return Abacus models
+        assert llm_provider._openai_client is not None, "OpenAI client should be initialized for Abacus"
+        with patch.object(llm_provider._openai_client, 'request_completions', new_callable=AsyncMock, return_value="Test Chat Name") as mock_abacus_request, \
+             patch("mito_ai.provider_manager.get_available_models", return_value=["Abacus/gpt-4.1", "Abacus/claude-haiku-4-5-20251001"]):
             result = await generate_short_chat_name(
                 user_message="What is the capital of France?",
                 assistant_message="The capital of France is Paris.",
                 llm_provider=llm_provider
             )
+            # Verify that the OpenAI client's request_completions was called (Abacus uses OpenAIClient)
+            mock_abacus_request.assert_called_once()  # type: ignore
+            # As a double check, if we have used the correct client, then we must get the correct result
+            assert result == "Test Chat Name"
+            return
     else:  # OpenAI
         # For OpenAI, patch the instance's method since the client is created in __init__
         assert llm_provider._openai_client is not None, "OpenAI client should be initialized"

mito_ai/tests/open_ai_utils_test.py

@@ -104,17 +104,16 @@ def test_prepare_request_data_and_headers_null_message() -> None:
     with patch("mito_ai.utils.open_ai_utils.get_user_field") as mock_get_user_field:
         mock_get_user_field.side_effect = ["test@example.com", "user123"]
         
-        with patch("mito_ai.utils.open_ai_utils.check_mito_server_quota"):
-            data, _ = _prepare_request_data_and_headers(
-                last_message_content=None,
-                ai_completion_data={},
-                timeout=30,
-                max_retries=3,
-                message_type=MessageType.CHAT
-            )
-
-            # Verify empty string is used for null message
-            assert data["user_input"] == ""
+        data, _ = _prepare_request_data_and_headers(
+            last_message_content=None,
+            ai_completion_data={},
+            timeout=30,
+            max_retries=3,
+            message_type=MessageType.CHAT
+        )
+
+        # Verify empty string is used for null message
+        assert data["user_input"] == ""
 
 def test_prepare_request_data_and_headers_caches_user_info() -> None:
     """Test that user info is cached after first call"""
@@ -125,28 +124,27 @@ def test_prepare_request_data_and_headers_caches_user_info() -> None:
         
         mock_get_user_field.side_effect = ["test@example.com", "user123"]
         
-        with patch("mito_ai.utils.open_ai_utils.check_mito_server_quota"):
-            # First call
-            data1, _ = _prepare_request_data_and_headers(
-                last_message_content="test",
-                ai_completion_data={},
-                timeout=30,
-                max_retries=3,
-                message_type=MessageType.CHAT
-            )
-
-            # Second call
-            data2, _ = _prepare_request_data_and_headers(
-                last_message_content="test",
-                ai_completion_data={},
-                timeout=30,
-                max_retries=3,
-                message_type=MessageType.CHAT
-            )
-
-            # Verify get_user_field was only called twice (once for email, once for user_id)
-            assert mock_get_user_field.call_count == 2
-            
-            # Verify both calls return same user info
-            assert data1["email"] == data2["email"] == "test@example.com"
-            assert data1["user_id"] == data2["user_id"] == "user123"
+        # First call
+        data1, _ = _prepare_request_data_and_headers(
+            last_message_content="test",
+            ai_completion_data={},
+            timeout=30,
+            max_retries=3,
+            message_type=MessageType.CHAT
+        )
+
+        # Second call
+        data2, _ = _prepare_request_data_and_headers(
+            last_message_content="test",
+            ai_completion_data={},
+            timeout=30,
+            max_retries=3,
+            message_type=MessageType.CHAT
+        )
+
+        # Verify get_user_field was only called twice (once for email, once for user_id)
+        assert mock_get_user_field.call_count == 2
+        
+        # Verify both calls return same user info
+        assert data1["email"] == data2["email"] == "test@example.com"
+        assert data1["user_id"] == data2["user_id"] == "user123"

mito_ai/tests/providers/test_azure.py

@@ -176,11 +176,11 @@ class TestAzureOpenAIClientCreation:
             openai_client = OpenAIClient(config=provider_config)
             
             # Test with gpt-4.1 model
-            resolved_model = openai_client._adjust_model_for_azure_or_ollama("gpt-4.1")
+            resolved_model = openai_client._adjust_model_for_provider("gpt-4.1")
             assert resolved_model == FAKE_AZURE_MODEL
             
             # Test with any other model
-            resolved_model = openai_client._adjust_model_for_azure_or_ollama("gpt-3.5-turbo")
+            resolved_model = openai_client._adjust_model_for_provider("gpt-3.5-turbo")
             assert resolved_model == FAKE_AZURE_MODEL
 
 

mito_ai/tests/test_constants.py

@@ -7,6 +7,7 @@ from mito_ai.constants import (
     ACTIVE_BASE_URL, MITO_PROD_BASE_URL, MITO_DEV_BASE_URL,
     MITO_STREAMLIT_DEV_BASE_URL, MITO_STREAMLIT_TEST_BASE_URL, ACTIVE_STREAMLIT_BASE_URL,
     COGNITO_CONFIG_DEV, ACTIVE_COGNITO_CONFIG,
+    parse_comma_separated_models,
 )
 
 
@@ -45,3 +46,92 @@ def test_cognito_config() -> Any:
 
     assert COGNITO_CONFIG_DEV == expected_config
     assert ACTIVE_COGNITO_CONFIG == COGNITO_CONFIG_DEV
+
+
+class TestParseCommaSeparatedModels:
+    """Tests for parse_comma_separated_models helper function."""
+    
+    def test_parse_models_no_quotes(self) -> None:
+        """Test parsing models without quotes."""
+        models_str = "litellm/openai/gpt-4o,litellm/anthropic/claude-3-5-sonnet"
+        result = parse_comma_separated_models(models_str)
+        assert result == ["litellm/openai/gpt-4o", "litellm/anthropic/claude-3-5-sonnet"]
+    
+    def test_parse_models_double_quotes(self) -> None:
+        """Test parsing models with double quotes."""
+        # Entire string quoted
+        models_str = '"litellm/openai/gpt-4o,litellm/anthropic/claude-3-5-sonnet"'
+        result = parse_comma_separated_models(models_str)
+        assert result == ["litellm/openai/gpt-4o", "litellm/anthropic/claude-3-5-sonnet"]
+        
+        # Individual models quoted
+        models_str = '"litellm/openai/gpt-4o","litellm/anthropic/claude-3-5-sonnet"'
+        result = parse_comma_separated_models(models_str)
+        assert result == ["litellm/openai/gpt-4o", "litellm/anthropic/claude-3-5-sonnet"]
+    
+    def test_parse_models_single_quotes(self) -> None:
+        """Test parsing models with single quotes."""
+        # Entire string quoted
+        models_str = "'litellm/openai/gpt-4o,litellm/anthropic/claude-3-5-sonnet'"
+        result = parse_comma_separated_models(models_str)
+        assert result == ["litellm/openai/gpt-4o", "litellm/anthropic/claude-3-5-sonnet"]
+        
+        # Individual models quoted
+        models_str = "'litellm/openai/gpt-4o','litellm/anthropic/claude-3-5-sonnet'"
+        result = parse_comma_separated_models(models_str)
+        assert result == ["litellm/openai/gpt-4o", "litellm/anthropic/claude-3-5-sonnet"]
+    
+    def test_parse_models_mixed_quotes(self) -> None:
+        """Test parsing models where some have single quotes and some have double quotes."""
+        # Some models with single quotes, some with double quotes
+        models_str = "'litellm/openai/gpt-4o',\"litellm/anthropic/claude-3-5-sonnet\""
+        result = parse_comma_separated_models(models_str)
+        # Should strip both types of quotes
+        assert result == ["litellm/openai/gpt-4o", "litellm/anthropic/claude-3-5-sonnet"]
+    
+    def test_parse_models_with_whitespace(self) -> None:
+        """Test parsing models with whitespace around commas and model names."""
+        models_str = " litellm/openai/gpt-4o , litellm/anthropic/claude-3-5-sonnet "
+        result = parse_comma_separated_models(models_str)
+        assert result == ["litellm/openai/gpt-4o", "litellm/anthropic/claude-3-5-sonnet"]
+    
+    def test_parse_models_empty_string(self) -> None:
+        """Test parsing empty string."""
+        result = parse_comma_separated_models("")
+        assert result == []
+    
+    def test_parse_models_single_model(self) -> None:
+        """Test parsing single model."""
+        models_str = "litellm/openai/gpt-4o"
+        result = parse_comma_separated_models(models_str)
+        assert result == ["litellm/openai/gpt-4o"]
+        
+        # With quotes
+        models_str = '"litellm/openai/gpt-4o"'
+        result = parse_comma_separated_models(models_str)
+        assert result == ["litellm/openai/gpt-4o"]
+    
+    def test_parse_models_abacus_format(self) -> None:
+        """Test parsing Abacus model format."""
+        models_str = "Abacus/gpt-4.1,Abacus/claude-haiku-4-5-20251001"
+        result = parse_comma_separated_models(models_str)
+        assert result == ["Abacus/gpt-4.1", "Abacus/claude-haiku-4-5-20251001"]
+        
+        # With quotes
+        models_str = '"Abacus/gpt-4.1","Abacus/claude-haiku-4-5-20251001"'
+        result = parse_comma_separated_models(models_str)
+        assert result == ["Abacus/gpt-4.1", "Abacus/claude-haiku-4-5-20251001"]
+    
+    @pytest.mark.parametrize("models_str,description", [
+        ('"model1,model2"', 'Double quotes, no space after comma'),
+        ("'model1,model2'", 'Single quotes, no space after comma'),
+        ("model1,model2", 'No quotes, no space after comma'),
+        ('"model1, model2"', 'Double quotes, space after comma'),
+        ("'model1, model2'", 'Single quotes, space after comma'),
+        ("model1, model2", 'No quotes, space after comma'),
+    ])
+    def test_parse_models_all_scenarios(self, models_str: str, description: str) -> None:
+        """Test all specific scenarios: quotes with and without spaces after commas."""
+        expected = ["model1", "model2"]
+        result = parse_comma_separated_models(models_str)
+        assert result == expected, f"Failed for {description}: {repr(models_str)}"