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

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