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

mito_ai/tests/providers/test_gemini_client.py

@@ -6,7 +6,7 @@ import ast
 import inspect
 import requests
 from mito_ai.gemini_client import GeminiClient, get_gemini_system_prompt_and_messages
-from mito_ai.utils.gemini_utils import get_gemini_completion_function_params, FAST_GEMINI_MODEL
+from mito_ai.utils.gemini_utils import get_gemini_completion_function_params
 from google.genai.types import Part, GenerateContentResponse, Candidate, Content
 from mito_ai.completions.models import ResponseFormatInfo, AgentResponse
 from unittest.mock import MagicMock, patch
@@ -156,19 +156,20 @@ async def test_json_response_handling_with_multiple_parts():
     assert result == 'Here is the JSON: {"key": "value"} End of response' 
     
 CUSTOM_MODEL = "smart-gemini-model"
-@pytest.mark.parametrize("message_type, expected_model", [
-    (MessageType.CHAT, CUSTOM_MODEL),  #
-    (MessageType.SMART_DEBUG, CUSTOM_MODEL),  #
-    (MessageType.CODE_EXPLAIN, CUSTOM_MODEL),  #
-    (MessageType.AGENT_EXECUTION, CUSTOM_MODEL),  #
-    (MessageType.AGENT_AUTO_ERROR_FIXUP, CUSTOM_MODEL),  #
-    (MessageType.INLINE_COMPLETION, FAST_GEMINI_MODEL),  #
-    (MessageType.CHAT_NAME_GENERATION, FAST_GEMINI_MODEL),  #
+@pytest.mark.parametrize("message_type", [
+    MessageType.CHAT,
+    MessageType.SMART_DEBUG,
+    MessageType.CODE_EXPLAIN,
+    MessageType.AGENT_EXECUTION,
+    MessageType.AGENT_AUTO_ERROR_FIXUP,
+    MessageType.INLINE_COMPLETION,
+    MessageType.CHAT_NAME_GENERATION,
 ])
 @pytest.mark.asyncio 
-async def test_get_completion_model_selection_based_on_message_type(message_type, expected_model):
+async def test_get_completion_model_selection_uses_passed_model(message_type):
     """
-    Tests that the correct model is selected based on the message type.
+    Tests that the model passed to the client is used as-is.
+    Model selection based on message type is now handled by ProviderManager.
     """
     with patch('google.genai.Client') as mock_genai_class:
         mock_client = MagicMock()
@@ -189,7 +190,7 @@ async def test_get_completion_model_selection_based_on_message_type(message_type
             response_format_info=None
         )
         
-        # Verify that generate_content was called with the expected model
+        # Verify that generate_content was called with the model that was passed (not overridden)
         mock_models.generate_content.assert_called_once()
         call_args = mock_models.generate_content.call_args
-        assert call_args[1]['model'] == expected_model 
+        assert call_args[1]['model'] == CUSTOM_MODEL 

mito_ai/tests/providers/test_model_resolution.py

@@ -6,19 +6,17 @@ These tests ensure that the correct model is chosen for each message type, for e
 """
 
 import pytest
-from mito_ai.utils.provider_utils import does_message_require_fast_model
+from mito_ai.utils.model_utils import get_fast_model_for_selected_model, get_smartest_model_for_selected_model
 from mito_ai.completions.models import MessageType
 from unittest.mock import AsyncMock, MagicMock, patch
-from mito_ai.completions.providers import OpenAIProvider
-from mito_ai.completions.models import MessageType
-from mito_ai.utils.provider_utils import does_message_require_fast_model
+from mito_ai.provider_manager import ProviderManager
 from traitlets.config import Config
 
 @pytest.fixture
 def provider_config() -> Config:
-    """Create a proper Config object for the OpenAIProvider."""
+    """Create a proper Config object for the ProviderManager."""
     config = Config()
-    config.OpenAIProvider = Config()
+    config.ProviderManager = Config()
     config.OpenAIClient = Config()
     return config
 
@@ -27,104 +25,162 @@ def mock_messages():
     """Sample messages for testing."""
     return [{"role": "user", "content": "Test message"}]
 
-# Test cases for different message types and their expected fast model requirement
-MESSAGE_TYPE_TEST_CASES = [
-    (MessageType.CHAT, False),
-    (MessageType.SMART_DEBUG, False),
-    (MessageType.CODE_EXPLAIN, False),
-    (MessageType.AGENT_EXECUTION, False),
-    (MessageType.AGENT_AUTO_ERROR_FIXUP, False),
-    (MessageType.INLINE_COMPLETION, True),
-    (MessageType.CHAT_NAME_GENERATION, True),
-]
-@pytest.mark.parametrize("message_type,expected_result", MESSAGE_TYPE_TEST_CASES)
-def test_does_message_require_fast_model(message_type: MessageType, expected_result: bool) -> None:
-    """Test that does_message_require_fast_model returns the correct boolean for each message type."""
-    assert does_message_require_fast_model(message_type) == expected_result
+@pytest.mark.asyncio
+async def test_request_completions_uses_fast_model_when_requested(provider_config: Config, mock_messages, monkeypatch: pytest.MonkeyPatch):
+    """Test that request_completions uses the correct model when use_fast_model=True."""
+    # Set up environment variables to ensure OpenAI provider is used
+    monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
+    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", "fake-key")
+    
+    # Mock the OpenAI API call instead of the entire client
+    mock_response = MagicMock()
+    mock_response.choices = [MagicMock()]
+    mock_response.choices[0].message.content = "Test Completion"
+    
+    with patch('openai.AsyncOpenAI') as mock_openai_class:
+        mock_openai_client = MagicMock()
+        mock_openai_client.chat.completions.create = AsyncMock(return_value=mock_response)
+        mock_openai_client.is_closed.return_value = False
+        mock_openai_class.return_value = mock_openai_client
+        
+        # Mock the validation that happens in OpenAIClient constructor
+        with patch('openai.OpenAI') as mock_sync_openai_class:
+            mock_sync_client = MagicMock()
+            mock_sync_client.models.list.return_value = MagicMock()
+            mock_sync_openai_class.return_value = mock_sync_client
+            
+            provider = ProviderManager(config=provider_config)
+            provider.set_selected_model("gpt-5.2")
+            await provider.request_completions(
+                message_type=MessageType.CHAT,
+                messages=mock_messages,
+                use_fast_model=True
+            )
+            
+            # Verify the model passed to the API call is the fast model
+            call_args = mock_openai_client.chat.completions.create.call_args
+            assert call_args[1]['model'] == get_fast_model_for_selected_model(provider.get_selected_model())
+
+@pytest.mark.asyncio
+async def test_stream_completions_uses_fast_model_when_requested(provider_config: Config, mock_messages, monkeypatch: pytest.MonkeyPatch):
+    """Test that stream_completions uses the correct model when use_fast_model=True."""
+    # Set up environment variables to ensure OpenAI provider is used
+    monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
+    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", "fake-key")
     
-def test_does_message_require_fast_model_raises_error_for_unknown_message_type():
-    """Test that does_message_require_fast_model raises an error for an unknown message type."""
-    with pytest.raises(ValueError):
-        does_message_require_fast_model('unknown_message_type') # type: ignore
+    # Mock the OpenAI API call instead of the entire client
+    mock_response = MagicMock()
+    mock_response.choices = [MagicMock()]
+    mock_response.choices[0].delta.content = "Test Stream Completion"
+    mock_response.choices[0].finish_reason = "stop"
+    
+    with patch('openai.AsyncOpenAI') as mock_openai_class:
+        mock_openai_client = MagicMock()
+        # Create an async generator for streaming
+        async def mock_stream():
+            yield mock_response
+        
+        mock_openai_client.chat.completions.create = AsyncMock(return_value=mock_stream())
+        mock_openai_client.is_closed.return_value = False
+        mock_openai_class.return_value = mock_openai_client
+        
+        # Mock the validation that happens in OpenAIClient constructor
+        with patch('openai.OpenAI') as mock_sync_openai_class:
+            mock_sync_client = MagicMock()
+            mock_sync_client.models.list.return_value = MagicMock()
+            mock_sync_openai_class.return_value = mock_sync_client
+            
+            provider = ProviderManager(config=provider_config)
+            provider.set_selected_model("gpt-5.2")
+            await provider.stream_completions(
+                message_type=MessageType.CHAT,
+                messages=mock_messages,
+                message_id="test_id",
+                thread_id="test_thread",
+                reply_fn=lambda x: None,
+                use_fast_model=True
+            )
+            
+            # Verify the model passed to the API call is the fast model
+            call_args = mock_openai_client.chat.completions.create.call_args
+            assert call_args[1]['model'] == get_fast_model_for_selected_model(provider.get_selected_model())
 
 @pytest.mark.asyncio
-async def test_request_completions_calls_does_message_require_fast_model(provider_config: Config, mock_messages, monkeypatch: pytest.MonkeyPatch):
-    """Test that request_completions calls does_message_require_fast_model and uses the correct model."""
+async def test_request_completions_uses_smartest_model_when_requested(provider_config: Config, mock_messages, monkeypatch: pytest.MonkeyPatch):
+    """Test that request_completions uses the correct model when use_smartest_model=True."""
     # Set up environment variables to ensure OpenAI provider is used
     monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
     monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", "fake-key")
     
-    with patch('mito_ai.utils.open_ai_utils.does_message_require_fast_model', wraps=does_message_require_fast_model) as mock_does_message_require_fast_model:
-        # Mock the OpenAI API call instead of the entire client
-        mock_response = MagicMock()
-        mock_response.choices = [MagicMock()]
-        mock_response.choices[0].message.content = "Test Completion"
+    # Mock the OpenAI API call instead of the entire client
+    mock_response = MagicMock()
+    mock_response.choices = [MagicMock()]
+    mock_response.choices[0].message.content = "Test Completion"
+    
+    with patch('openai.AsyncOpenAI') as mock_openai_class:
+        mock_openai_client = MagicMock()
+        mock_openai_client.chat.completions.create = AsyncMock(return_value=mock_response)
+        mock_openai_client.is_closed.return_value = False
+        mock_openai_class.return_value = mock_openai_client
         
-        with patch('openai.AsyncOpenAI') as mock_openai_class:
-            mock_openai_client = MagicMock()
-            mock_openai_client.chat.completions.create = AsyncMock(return_value=mock_response)
-            mock_openai_client.is_closed.return_value = False
-            mock_openai_class.return_value = mock_openai_client
+        # Mock the validation that happens in OpenAIClient constructor
+        with patch('openai.OpenAI') as mock_sync_openai_class:
+            mock_sync_client = MagicMock()
+            mock_sync_client.models.list.return_value = MagicMock()
+            mock_sync_openai_class.return_value = mock_sync_client
             
-            # Mock the validation that happens in OpenAIClient constructor
-            with patch('openai.OpenAI') as mock_sync_openai_class:
-                mock_sync_client = MagicMock()
-                mock_sync_client.models.list.return_value = MagicMock()
-                mock_sync_openai_class.return_value = mock_sync_client
-                
-                provider = OpenAIProvider(config=provider_config)
-                await provider.request_completions(
-                    message_type=MessageType.CHAT,
-                    messages=mock_messages,
-                    model="gpt-3.5",
-                )
-                
-                mock_does_message_require_fast_model.assert_called_once_with(MessageType.CHAT)
-                # Verify the model passed to the API call
-                call_args = mock_openai_client.chat.completions.create.call_args
-                assert call_args[1]['model'] == "gpt-3.5"
+            provider = ProviderManager(config=provider_config)
+            provider.set_selected_model("gpt-4.1")
+            await provider.request_completions(
+                message_type=MessageType.CHAT,
+                messages=mock_messages,
+                use_smartest_model=True
+            )
+            
+            # Verify the model passed to the API call is the smartest model
+            call_args = mock_openai_client.chat.completions.create.call_args
+            assert call_args[1]['model'] == get_smartest_model_for_selected_model(provider.get_selected_model())
 
 @pytest.mark.asyncio
-async def test_stream_completions_calls_does_message_require_fast_model(provider_config: Config, mock_messages, monkeypatch: pytest.MonkeyPatch):
-    """Test that stream_completions calls does_message_require_fast_model and uses the correct model."""
+async def test_stream_completions_uses_smartest_model_when_requested(provider_config: Config, mock_messages, monkeypatch: pytest.MonkeyPatch):
+    """Test that stream_completions uses the correct model when use_smartest_model=True."""
     # Set up environment variables to ensure OpenAI provider is used
     monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
     monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", "fake-key")
     
-    with patch('mito_ai.utils.open_ai_utils.does_message_require_fast_model', wraps=does_message_require_fast_model) as mock_does_message_require_fast_model:
-        # Mock the OpenAI API call instead of the entire client
-        mock_response = MagicMock()
-        mock_response.choices = [MagicMock()]
-        mock_response.choices[0].delta.content = "Test Stream Completion"
-        mock_response.choices[0].finish_reason = "stop"
+    # Mock the OpenAI API call instead of the entire client
+    mock_response = MagicMock()
+    mock_response.choices = [MagicMock()]
+    mock_response.choices[0].delta.content = "Test Stream Completion"
+    mock_response.choices[0].finish_reason = "stop"
+    
+    with patch('openai.AsyncOpenAI') as mock_openai_class:
+        mock_openai_client = MagicMock()
+        # Create an async generator for streaming
+        async def mock_stream():
+            yield mock_response
+        
+        mock_openai_client.chat.completions.create = AsyncMock(return_value=mock_stream())
+        mock_openai_client.is_closed.return_value = False
+        mock_openai_class.return_value = mock_openai_client
         
-        with patch('openai.AsyncOpenAI') as mock_openai_class:
-            mock_openai_client = MagicMock()
-            # Create an async generator for streaming
-            async def mock_stream():
-                yield mock_response
+        # Mock the validation that happens in OpenAIClient constructor
+        with patch('openai.OpenAI') as mock_sync_openai_class:
+            mock_sync_client = MagicMock()
+            mock_sync_client.models.list.return_value = MagicMock()
+            mock_sync_openai_class.return_value = mock_sync_client
             
-            mock_openai_client.chat.completions.create = AsyncMock(return_value=mock_stream())
-            mock_openai_client.is_closed.return_value = False
-            mock_openai_class.return_value = mock_openai_client
+            provider = ProviderManager(config=provider_config)
+            provider.set_selected_model("gpt-4.1")
+            await provider.stream_completions(
+                message_type=MessageType.CHAT,
+                messages=mock_messages,
+                message_id="test_id",
+                thread_id="test_thread",
+                reply_fn=lambda x: None,
+                use_smartest_model=True
+            )
             
-            # Mock the validation that happens in OpenAIClient constructor
-            with patch('openai.OpenAI') as mock_sync_openai_class:
-                mock_sync_client = MagicMock()
-                mock_sync_client.models.list.return_value = MagicMock()
-                mock_sync_openai_class.return_value = mock_sync_client
-                
-                provider = OpenAIProvider(config=provider_config)
-                await provider.stream_completions(
-                    message_type=MessageType.CHAT,
-                    messages=mock_messages,
-                    model="gpt-3.5",
-                    message_id="test_id",
-                    thread_id="test_thread",
-                    reply_fn=lambda x: None
-                )
-                
-                mock_does_message_require_fast_model.assert_called_once_with(MessageType.CHAT)
-                # Verify the model passed to the API call
-                call_args = mock_openai_client.chat.completions.create.call_args
-                assert call_args[1]['model'] == "gpt-3.5"
+            # Verify the model passed to the API call is the smartest model
+            call_args = mock_openai_client.chat.completions.create.call_args
+            assert call_args[1]['model'] == get_smartest_model_for_selected_model(provider.get_selected_model())

mito_ai/tests/providers/test_openai_client.py

@@ -3,25 +3,25 @@
 
 import pytest
 from mito_ai.openai_client import OpenAIClient
-from mito_ai.utils.open_ai_utils import FAST_OPENAI_MODEL
 from mito_ai.completions.models import MessageType
 from unittest.mock import MagicMock, patch, AsyncMock
 from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
 
 CUSTOM_MODEL = "smart-openai-model"
-@pytest.mark.parametrize("message_type, expected_model", [
-    (MessageType.CHAT, CUSTOM_MODEL),  #
-    (MessageType.SMART_DEBUG, CUSTOM_MODEL),  #
-    (MessageType.CODE_EXPLAIN, CUSTOM_MODEL),  #
-    (MessageType.AGENT_EXECUTION, CUSTOM_MODEL),  #
-    (MessageType.AGENT_AUTO_ERROR_FIXUP, CUSTOM_MODEL),  #
-    (MessageType.INLINE_COMPLETION, FAST_OPENAI_MODEL),  #
-    (MessageType.CHAT_NAME_GENERATION, FAST_OPENAI_MODEL),  #
+@pytest.mark.parametrize("message_type", [
+    MessageType.CHAT,
+    MessageType.SMART_DEBUG,
+    MessageType.CODE_EXPLAIN,
+    MessageType.AGENT_EXECUTION,
+    MessageType.AGENT_AUTO_ERROR_FIXUP,
+    MessageType.INLINE_COMPLETION,
+    MessageType.CHAT_NAME_GENERATION,
 ])
 @pytest.mark.asyncio 
-async def test_model_selection_based_on_message_type(message_type, expected_model):
+async def test_model_selection_uses_passed_model(message_type):
     """
-    Tests that the correct model is selected based on the message type.
+    Tests that the model passed to the client is used as-is.
+    Model selection based on message type is now handled by ProviderManager.
     """
     client = OpenAIClient(api_key="test_key") # type: ignore
     
@@ -51,7 +51,203 @@ async def test_model_selection_based_on_message_type(message_type, expected_mode
             response_format_info=None
         )
         
-        # Verify that create was called with the expected model
+        # Verify that create was called with the model that was passed (not overridden)
         mock_create.assert_called_once()
         call_args = mock_create.call_args
-        assert call_args[1]['model'] == expected_model
+        assert call_args[1]['model'] == CUSTOM_MODEL
+
+@pytest.mark.asyncio
+async def test_openai_client_uses_fast_model_from_provider_manager_without_override():
+    """Test that OpenAI client uses the fast model passed from ProviderManager without internal override."""
+    from mito_ai.utils.model_utils import get_fast_model_for_selected_model
+    
+    client = OpenAIClient(api_key="test_key") # type: ignore
+    
+    # Mock the _build_openai_client method to return our mock client
+    with patch.object(client, '_build_openai_client') as mock_build_client, \
+         patch('openai.AsyncOpenAI') as mock_openai_class:
+        
+        mock_client = MagicMock()
+        mock_chat = MagicMock()
+        mock_completions = MagicMock()
+        mock_client.chat = mock_chat
+        mock_chat.completions = mock_completions
+        mock_openai_class.return_value = mock_client
+        mock_build_client.return_value = mock_client
+        
+        # Create an async mock for the create method
+        mock_create = AsyncMock()
+        mock_create.return_value = MagicMock(
+            choices=[MagicMock(message=MagicMock(content="test"))]
+        )
+        mock_completions.create = mock_create
+
+        # Use a fast model that would be selected by ProviderManager
+        fast_model = get_fast_model_for_selected_model("gpt-5.2")
+        
+        await client.request_completions(
+            message_type=MessageType.CHAT,
+            messages=[{"role": "user", "content": "Test message"}],
+            model=fast_model,
+            response_format_info=None
+        )
+        
+        # Verify that create was called with the fast model that was passed (not overridden)
+        mock_create.assert_called_once()
+        call_args = mock_create.call_args
+        assert call_args[1]['model'] == fast_model
+
+@pytest.mark.asyncio
+async def test_openai_client_uses_smartest_model_from_provider_manager_without_override():
+    """Test that OpenAI client uses the smartest model passed from ProviderManager without internal override."""
+    from mito_ai.utils.model_utils import get_smartest_model_for_selected_model
+    
+    client = OpenAIClient(api_key="test_key") # type: ignore
+    
+    # Mock the _build_openai_client method to return our mock client
+    with patch.object(client, '_build_openai_client') as mock_build_client, \
+         patch('openai.AsyncOpenAI') as mock_openai_class:
+        
+        mock_client = MagicMock()
+        mock_chat = MagicMock()
+        mock_completions = MagicMock()
+        mock_client.chat = mock_chat
+        mock_chat.completions = mock_completions
+        mock_openai_class.return_value = mock_client
+        mock_build_client.return_value = mock_client
+        
+        # Create an async mock for the create method
+        mock_create = AsyncMock()
+        mock_create.return_value = MagicMock(
+            choices=[MagicMock(message=MagicMock(content="test"))]
+        )
+        mock_completions.create = mock_create
+
+        # Use a smartest model that would be selected by ProviderManager
+        smartest_model = get_smartest_model_for_selected_model("gpt-4.1")
+        
+        await client.request_completions(
+            message_type=MessageType.CHAT,
+            messages=[{"role": "user", "content": "Test message"}],
+            model=smartest_model,
+            response_format_info=None
+        )
+        
+        # Verify that create was called with the smartest model that was passed (not overridden)
+        mock_create.assert_called_once()
+        call_args = mock_create.call_args
+        assert call_args[1]['model'] == smartest_model
+
+@pytest.mark.asyncio
+async def test_openai_client_stream_uses_fast_model_from_provider_manager_without_override():
+    """Test that OpenAI client stream_completions uses the fast model passed from ProviderManager without internal override."""
+    from mito_ai.utils.model_utils import get_fast_model_for_selected_model
+    
+    client = OpenAIClient(api_key="test_key") # type: ignore
+    
+    # Mock the _build_openai_client method to return our mock client
+    with patch.object(client, '_build_openai_client') as mock_build_client, \
+         patch('openai.AsyncOpenAI') as mock_openai_class:
+        
+        mock_client = MagicMock()
+        mock_chat = MagicMock()
+        mock_completions = MagicMock()
+        mock_client.chat = mock_chat
+        mock_chat.completions = mock_completions
+        mock_openai_class.return_value = mock_client
+        mock_build_client.return_value = mock_client
+        
+        # Create an async generator for streaming
+        async def mock_stream():
+            mock_chunk = MagicMock()
+            mock_chunk.choices = [MagicMock()]
+            mock_chunk.choices[0].delta.content = "test"
+            mock_chunk.choices[0].finish_reason = None
+            yield mock_chunk
+            mock_final_chunk = MagicMock()
+            mock_final_chunk.choices = [MagicMock()]
+            mock_final_chunk.choices[0].delta.content = ""
+            mock_final_chunk.choices[0].finish_reason = "stop"
+            yield mock_final_chunk
+        
+        mock_create = AsyncMock(return_value=mock_stream())
+        mock_completions.create = mock_create
+
+        # Use a fast model that would be selected by ProviderManager
+        fast_model = get_fast_model_for_selected_model("gpt-5.2")
+        
+        reply_chunks = []
+        def mock_reply(chunk):
+            reply_chunks.append(chunk)
+        
+        await client.stream_completions(
+            message_type=MessageType.CHAT,
+            messages=[{"role": "user", "content": "Test message"}],
+            model=fast_model,
+            message_id="test-id",
+            thread_id="test-thread",
+            reply_fn=mock_reply,
+            response_format_info=None
+        )
+        
+        # Verify that create was called with the fast model that was passed (not overridden)
+        mock_create.assert_called_once()
+        call_args = mock_create.call_args
+        assert call_args[1]['model'] == fast_model
+
+@pytest.mark.asyncio
+async def test_openai_client_stream_uses_smartest_model_from_provider_manager_without_override():
+    """Test that OpenAI client stream_completions uses the smartest model passed from ProviderManager without internal override."""
+    from mito_ai.utils.model_utils import get_smartest_model_for_selected_model
+    
+    client = OpenAIClient(api_key="test_key") # type: ignore
+    
+    # Mock the _build_openai_client method to return our mock client
+    with patch.object(client, '_build_openai_client') as mock_build_client, \
+         patch('openai.AsyncOpenAI') as mock_openai_class:
+        
+        mock_client = MagicMock()
+        mock_chat = MagicMock()
+        mock_completions = MagicMock()
+        mock_client.chat = mock_chat
+        mock_chat.completions = mock_completions
+        mock_openai_class.return_value = mock_client
+        mock_build_client.return_value = mock_client
+        
+        # Create an async generator for streaming
+        async def mock_stream():
+            mock_chunk = MagicMock()
+            mock_chunk.choices = [MagicMock()]
+            mock_chunk.choices[0].delta.content = "test"
+            mock_chunk.choices[0].finish_reason = None
+            yield mock_chunk
+            mock_final_chunk = MagicMock()
+            mock_final_chunk.choices = [MagicMock()]
+            mock_final_chunk.choices[0].delta.content = ""
+            mock_final_chunk.choices[0].finish_reason = "stop"
+            yield mock_final_chunk
+        
+        mock_create = AsyncMock(return_value=mock_stream())
+        mock_completions.create = mock_create
+
+        # Use a smartest model that would be selected by ProviderManager
+        smartest_model = get_smartest_model_for_selected_model("gpt-4.1")
+        
+        reply_chunks = []
+        def mock_reply(chunk):
+            reply_chunks.append(chunk)
+        
+        await client.stream_completions(
+            message_type=MessageType.CHAT,
+            messages=[{"role": "user", "content": "Test message"}],
+            model=smartest_model,
+            message_id="test-id",
+            thread_id="test-thread",
+            reply_fn=mock_reply,
+            response_format_info=None
+        )
+        
+        # Verify that create was called with the smartest model that was passed (not overridden)
+        mock_create.assert_called_once()
+        call_args = mock_create.call_args
+        assert call_args[1]['model'] == smartest_model

mito_ai/tests/providers/test_provider_limits.py

@@ -2,7 +2,7 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 import pytest
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 from mito_ai.tests.providers.utils import mock_openai_client, patch_server_limits
 from mito_ai.utils.server_limits import OS_MONTHLY_AI_COMPLETIONS_LIMIT
 from traitlets.config import Config
@@ -11,9 +11,9 @@ FAKE_API_KEY = "sk-1234567890"
 
 @pytest.fixture
 def provider_config() -> Config:
-    """Create a proper Config object for the OpenAIProvider."""
+    """Create a proper Config object for the ProviderManager."""
     config = Config()
-    config.OpenAIProvider = Config()
+    config.ProviderManager = Config()
     config.OpenAIClient = Config()
     return config
 
@@ -36,7 +36,7 @@ def test_openai_provider_with_limits(
         patch_server_limits(is_pro=is_pro, completion_count=completion_count),
         mock_openai_client()
     ):
-        llm = OpenAIProvider(config=provider_config)
+        llm = ProviderManager(config=provider_config)
         capabilities = llm.capabilities
-        assert "user key" in capabilities.provider
+        assert "OpenAI" in capabilities.provider
         assert llm.last_error is None