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

mito_ai/tests/providers/test_azure.py

@@ -9,7 +9,7 @@ import pytest
 from traitlets.config import Config
 from openai.types.chat import ChatCompletionMessageParam
 
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 from mito_ai.completions.models import (
     MessageType,
     AICapabilities,
@@ -29,9 +29,9 @@ FAKE_AZURE_API_VERSION = "2024-12-01-preview"
 
 @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
 
@@ -40,7 +40,7 @@ def provider_config() -> Config:
 def reset_env_vars(monkeypatch: pytest.MonkeyPatch) -> None:
     """Reset all environment variables before each test."""
     for var in [
-        "OPENAI_API_KEY", "CLAUDE_API_KEY", "GEMINI_API_KEY", "OLLAMA_MODEL",
+        "OPENAI_API_KEY", "ANTHROPIC_API_KEY", "GEMINI_API_KEY", "OLLAMA_MODEL",
         "AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_MODEL", 
         "AZURE_OPENAI_API_VERSION"
     ]:
@@ -405,7 +405,7 @@ class TestAzureOpenAIStreamCompletions:
 
 
 class TestAzureOpenAIProviderIntegration:
-    """Test Azure OpenAI integration through the OpenAIProvider."""
+    """Test Azure OpenAI integration through the ProviderManager."""
     
     @pytest.mark.asyncio
     @pytest.mark.parametrize("message_type", COMPLETION_MESSAGE_TYPES)
@@ -415,7 +415,7 @@ class TestAzureOpenAIProviderIntegration:
         provider_config: Config,
         message_type: MessageType
     ) -> None:
-        """Test that OpenAIProvider uses Azure OpenAI when gpt-4.1 is requested and Azure is configured."""
+        """Test that ProviderManager uses Azure OpenAI when gpt-4.1 is requested and Azure is configured."""
         
         # Mock the response
         mock_response = MagicMock()
@@ -428,7 +428,8 @@ class TestAzureOpenAIProviderIntegration:
             mock_azure_client.is_closed.return_value = False
             mock_azure_client_class.return_value = mock_azure_client
             
-            provider = OpenAIProvider(config=provider_config)
+            provider = ProviderManager(config=provider_config)
+            provider.set_selected_model("gpt-4.1")
             
             messages: List[ChatCompletionMessageParam] = [
                 {"role": "user", "content": "Test message"}
@@ -437,7 +438,6 @@ class TestAzureOpenAIProviderIntegration:
             completion = await provider.request_completions(
                 message_type=message_type,
                 messages=messages,
-                model="gpt-4.1"
             )
             
             # Verify the completion was returned
@@ -461,7 +461,7 @@ class TestAzureOpenAIProviderIntegration:
         provider_config: Config,
         message_type: MessageType
     ) -> None:
-        """Test that OpenAIProvider stream_completions uses Azure OpenAI when gpt-4.1 is requested and Azure is configured."""
+        """Test that ProviderManager stream_completions uses Azure OpenAI when gpt-4.1 is requested and Azure is configured."""
         
         # Mock the streaming response
         mock_chunk1 = MagicMock()
@@ -484,7 +484,8 @@ class TestAzureOpenAIProviderIntegration:
             mock_azure_client.is_closed.return_value = False
             mock_azure_client_class.return_value = mock_azure_client
             
-            provider = OpenAIProvider(config=provider_config)
+            provider = ProviderManager(config=provider_config)
+            provider.set_selected_model("gpt-4.1")
             
             messages: List[ChatCompletionMessageParam] = [
                 {"role": "user", "content": "Test message"}
@@ -497,7 +498,6 @@ class TestAzureOpenAIProviderIntegration:
             completion = await provider.stream_completions(
                 message_type=message_type,
                 messages=messages,
-                model="gpt-4.1",
                 message_id="test-id",
                 thread_id="test-thread",
                 reply_fn=mock_reply
@@ -554,8 +554,8 @@ class TestAzureOpenAIConfigurationPriority:
         """Test that Azure OpenAI is used even when Claude key is available."""
         
         # Set Claude key (this should be overridden by Azure OpenAI)
-        monkeypatch.setenv("CLAUDE_API_KEY", "claude-key")
-        monkeypatch.setattr("mito_ai.constants.CLAUDE_API_KEY", "claude-key")
+        monkeypatch.setenv("ANTHROPIC_API_KEY", "claude-key")
+        monkeypatch.setattr("mito_ai.constants.ANTHROPIC_API_KEY", "claude-key")
         
         with patch("openai.AsyncAzureOpenAI") as mock_azure_client:
             openai_client = OpenAIClient(config=provider_config)

mito_ai/tests/providers/test_capabilities.py

@@ -3,7 +3,7 @@
 
 import pytest
 from unittest.mock import MagicMock, patch
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 from mito_ai.tests.providers.utils import mock_azure_openai_client, mock_openai_client, patch_server_limits
 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
 
@@ -22,7 +22,7 @@ def provider_config() -> Config:
         "name": "mito_server_fallback_no_keys",
         "setup": {
             "OPENAI_API_KEY": None,
-            "CLAUDE_API_KEY": None, 
+            "ANTHROPIC_API_KEY": None, 
             "GEMINI_API_KEY": None,
             "is_azure_configured": False,
         },
@@ -33,45 +33,45 @@ def provider_config() -> Config:
         "name": "claude_when_only_claude_key",
         "setup": {
             "OPENAI_API_KEY": None,
-            "CLAUDE_API_KEY": "claude-test-key",
+            "ANTHROPIC_API_KEY": "claude-test-key",
             "GEMINI_API_KEY": None,
             "is_azure_configured": False,
         },
         "expected_provider": "Claude",
-        "expected_key_type": "claude"
+        "expected_key_type": "user_key"
     },
     {
         "name": "gemini_when_only_gemini_key",
         "setup": {
             "OPENAI_API_KEY": None,
-            "CLAUDE_API_KEY": None,
+            "ANTHROPIC_API_KEY": None,
             "GEMINI_API_KEY": "gemini-test-key",
             "is_azure_configured": False,
         },
         "expected_provider": "Gemini", 
-        "expected_key_type": "gemini"
+        "expected_key_type": "user_key"
     },
     {
         "name": "openai_when_openai_key",
         "setup": {
             "OPENAI_API_KEY": 'openai-test-key',
-            "CLAUDE_API_KEY": None,
+            "ANTHROPIC_API_KEY": None,
             "GEMINI_API_KEY": None,
             "is_azure_configured": False,
         },
-        "expected_provider": "OpenAI (user key)",
+        "expected_provider": "OpenAI",
         "expected_key_type": "user_key"
     },
     {
         "name": "claude_priority_over_gemini",
         "setup": {
             "OPENAI_API_KEY": None,
-            "CLAUDE_API_KEY": "claude-test-key",
+            "ANTHROPIC_API_KEY": "claude-test-key",
             "GEMINI_API_KEY": "gemini-test-key",
             "is_azure_configured": False,
         },
         "expected_provider": "Claude",
-        "expected_key_type": "claude"
+        "expected_key_type": "user_key"
     },
 ])
 def test_provider_capabilities_real_logic(
@@ -79,7 +79,7 @@ def test_provider_capabilities_real_logic(
     monkeypatch: pytest.MonkeyPatch, 
     provider_config: Config
 ) -> None:
-    """Test the actual provider selection logic in OpenAIProvider.capabilities"""
+    """Test the actual provider selection logic in ProviderManager.capabilities"""
     
     # Set up the environment based on test case
     setup = test_case["setup"]
@@ -97,9 +97,6 @@ def test_provider_capabilities_real_logic(
         else:
             monkeypatch.setattr(f"mito_ai.constants.{key}", value)
     
-    # Clear the provider config API key to ensure it uses constants
-    provider_config.OpenAIProvider.api_key = None
-    
     # Mock HTTP calls but let the real logic run
     with patch("openai.OpenAI") as mock_openai_constructor:
         with patch("openai.AsyncOpenAI") as mock_async_openai:
@@ -112,7 +109,7 @@ def test_provider_capabilities_real_logic(
                 # Mock server limits for Mito server fallback
                 with patch_server_limits():
                     # NOW create the provider after ALL mocks are set up
-                    llm = OpenAIProvider(config=provider_config)
+                    llm = ProviderManager(config=provider_config)
                     
                     # Test capabilities 
                     capabilities = llm.capabilities

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())