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

mito_ai/streamlit_preview/handlers.py

@@ -11,15 +11,18 @@ from mito_ai.utils.telemetry_utils import log_streamlit_app_conversion_error, lo
 from mito_ai.completions.models import MessageType
 from mito_ai.utils.error_classes import StreamlitConversionError, StreamlitPreviewError
 from mito_ai.streamlit_conversion.streamlit_agent_handler import streamlit_handler
+from mito_ai.provider_manager import ProviderManager
 import traceback
 
 
 class StreamlitPreviewHandler(APIHandler):
     """REST handler for streamlit preview operations."""
 
-    def initialize(self) -> None:
+    def initialize(self, llm: ProviderManager) -> None:
         """Initialize the handler."""
+        super().initialize()
         self.preview_manager = StreamlitPreviewManager()
+        self._llm = llm
 
     @tornado.web.authenticated
     
@@ -45,11 +48,11 @@ class StreamlitPreviewHandler(APIHandler):
                     print("[Mito AI] Force recreating streamlit app")
 
                 # Create a new app 
-                await streamlit_handler(True, absolute_notebook_path, app_file_name, streamlit_app_prompt)
+                await streamlit_handler(True, absolute_notebook_path, app_file_name, streamlit_app_prompt, self._llm)
             elif streamlit_app_prompt != '':
                 # Update an existing app if there is a prompt provided. Otherwise, the user is just
                 # starting an existing app so we can skip the streamlit_handler all together
-                await streamlit_handler(False, absolute_notebook_path, app_file_name, streamlit_app_prompt)
+                await streamlit_handler(False, absolute_notebook_path, app_file_name, streamlit_app_prompt, self._llm)
 
             # Start preview
             # TODO: There's a bug here where when the user rebuilds and already running app. Instead of 

mito_ai/streamlit_preview/urls.py

@@ -4,12 +4,14 @@
 from typing import Any, List, Tuple
 from jupyter_server.utils import url_path_join
 from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
+from mito_ai.provider_manager import ProviderManager
 
-def get_streamlit_preview_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
+def get_streamlit_preview_urls(base_url: str, provider_manager: ProviderManager) -> List[Tuple[str, Any, dict]]:
     """Get all streamlit preview related URL patterns.
 
     Args:
         base_url: The base URL for the Jupyter server
+        provider_manager: The ProviderManager instance
 
     Returns:
         List of (url_pattern, handler_class, handler_kwargs) tuples
@@ -17,6 +19,6 @@ def get_streamlit_preview_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
     BASE_URL = base_url + "/mito-ai"
     
     return [
-        (url_path_join(BASE_URL, "streamlit-preview"), StreamlitPreviewHandler, {}),
-        (url_path_join(BASE_URL, "streamlit-preview/(.+)"), StreamlitPreviewHandler, {}),
+        (url_path_join(BASE_URL, "streamlit-preview"), StreamlitPreviewHandler, {"llm": provider_manager}),
+        (url_path_join(BASE_URL, "streamlit-preview/(.+)"), StreamlitPreviewHandler, {"llm": provider_manager}),
     ]

mito_ai/tests/message_history/test_generate_short_chat_name.py

@@ -5,24 +5,25 @@ import pytest
 from unittest.mock import AsyncMock, MagicMock, patch
 from traitlets.config import Config
 from mito_ai.completions.message_history import generate_short_chat_name
-from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.provider_manager import ProviderManager
 
 
 @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
 
 
 # Test cases for different models and their expected providers/fast models
 PROVIDER_TEST_CASES = [
-    # (model, client_patch_path)
-    ("gpt-4.1", "mito_ai.completions.providers.OpenAIClient"),
-    ("claude-3-5-sonnet-20241022", "mito_ai.completions.providers.AnthropicClient"),
-    ("gemini-2.0-flash-exp", "mito_ai.completions.providers.GeminiClient")
+    # (model, client_patch_path) - patch where the classes are used (in provider_manager)
+    ("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
 ]
 
 @pytest.mark.parametrize("selected_model,client_patch_path", PROVIDER_TEST_CASES)
@@ -37,31 +38,77 @@ async def test_generate_short_chat_name_uses_correct_provider_and_fast_model(
     
     # Set up environment variables for all providers
     monkeypatch.setenv("OPENAI_API_KEY", "fake-openai-key")
-    monkeypatch.setenv("CLAUDE_API_KEY", "fake-claude-key")
+    monkeypatch.setenv("ANTHROPIC_API_KEY", "fake-claude-key")
     monkeypatch.setenv("GEMINI_API_KEY", "fake-gemini-key")
     monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", "fake-openai-key")
-    monkeypatch.setattr("mito_ai.constants.CLAUDE_API_KEY", "fake-claude-key")
+    monkeypatch.setattr("mito_ai.constants.ANTHROPIC_API_KEY", "fake-claude-key")
     monkeypatch.setattr("mito_ai.constants.GEMINI_API_KEY", "fake-gemini-key")
     
+    # Set up LiteLLM constants if testing LiteLLM
+    if "LiteLLMClient" in client_patch_path:
+        # 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"])
+        # 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"])
+        # Mock is_enterprise to return True so LiteLLM 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")
     
+    # Create the ProviderManager first
+    llm_provider = ProviderManager(config=provider_config)
+    
+    # Set the selected model (this is required for the ProviderManager to use the correct model)
+    llm_provider.set_selected_model(selected_model)
+    
     # Patch the specific client class that should be used based on the model
-    # We need to patch before creating the OpenAIProvider since OpenAI client is created in constructor
-    with patch(client_patch_path, return_value=mock_client):
-        # Create the OpenAIProvider after patching so the mock client is used
-        llm_provider = OpenAIProvider(config=provider_config)
-        
-        # Test the function
-        result = await generate_short_chat_name(
-            user_message="What is the capital of France?",
-            assistant_message="The capital of France is Paris.",
-            model=selected_model,
-            llm_provider=llm_provider
-        )
-        
-    # Verify that the correct client's request_completions was called
+    # For Anthropic, Gemini, and LiteLLM, new instances are created in request_completions, so we patch the class
+    # For OpenAI, the instance is created in __init__, so we patch the instance's method
+    if "AnthropicClient" in client_patch_path:
+        with patch(client_patch_path, return_value=mock_client):
+            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 "GeminiClient" in client_patch_path:
+        with patch(client_patch_path, return_value=mock_client):
+            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 "LiteLLMClient" in client_patch_path:
+        # 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"]):
+            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
+            )
+    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"
+        with patch.object(llm_provider._openai_client, 'request_completions', new_callable=AsyncMock, return_value="Test Chat Name") as mock_openai_request:
+            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
+            mock_openai_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
+    
+    # Verify that the correct client's request_completions was called (for Anthropic, Gemini, and LiteLLM)
     mock_client.request_completions.assert_called_once()
 
     # As a double check, if we have used the correct client, then we must get the correct result
@@ -74,13 +121,12 @@ async def test_generate_short_chat_name_cleans_gemini_response() -> None:
     """Test that generate_short_chat_name properly cleans Gemini-style responses with quotes and newlines."""
     
     # Create mock llm_provider that returns a response with quotes and newlines
-    mock_llm_provider = MagicMock(spec=OpenAIProvider)
+    mock_llm_provider = MagicMock(spec=ProviderManager)
     mock_llm_provider.request_completions = AsyncMock(return_value='"France Geography Discussion\n"')
     
     result = await generate_short_chat_name(
         user_message="What is the capital of France?",
         assistant_message="The capital of France is Paris.",
-        model="gemini-2.0-flash-exp",
         llm_provider=mock_llm_provider
     )
     
@@ -95,13 +141,12 @@ async def test_generate_short_chat_name_handles_empty_response() -> None:
     """Test that generate_short_chat_name handles empty or None responses gracefully."""
     
     # Test with empty string response
-    mock_llm_provider = MagicMock(spec=OpenAIProvider)
+    mock_llm_provider = MagicMock(spec=ProviderManager)
     mock_llm_provider.request_completions = AsyncMock(return_value="")
     
     result = await generate_short_chat_name(
         user_message="Test message",
         assistant_message="Test response",
-        model="gpt-4.1",
         llm_provider=mock_llm_provider
     )
     
@@ -113,7 +158,6 @@ async def test_generate_short_chat_name_handles_empty_response() -> None:
     result = await generate_short_chat_name(
         user_message="Test message", 
         assistant_message="Test response",
-        model="gpt-4.1",
         llm_provider=mock_llm_provider
     )
     

mito_ai/tests/providers/test_anthropic_client.py

@@ -3,11 +3,10 @@
 
 import pytest
 from mito_ai.anthropic_client import get_anthropic_system_prompt_and_messages, get_anthropic_system_prompt_and_messages_with_caching, add_cache_control_to_message, extract_and_parse_anthropic_json_response, AnthropicClient
-from mito_ai.utils.anthropic_utils import FAST_ANTHROPIC_MODEL
 from anthropic.types import Message, TextBlock, ToolUseBlock, Usage, ToolUseBlock, Message, Usage, TextBlock
 from openai.types.chat import ChatCompletionMessageParam, ChatCompletionUserMessageParam, ChatCompletionAssistantMessageParam, ChatCompletionSystemMessageParam
 from mito_ai.completions.models import MessageType
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 import anthropic
 from typing import List, Dict, cast
 
@@ -233,24 +232,25 @@ def test_tool_use_without_agent_response():
     assert "No valid AgentResponse format found" in str(exc_info.value)
 
 CUSTOM_MODEL = "smart-anthropic-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_ANTHROPIC_MODEL),  #
-    (MessageType.CHAT_NAME_GENERATION, FAST_ANTHROPIC_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 = AnthropicClient(api_key="test_key")
     
-    # Mock the messages.create method directly
-    with patch.object(client.client.messages, 'create') as mock_create: # type: ignore
+    # Mock the beta.messages.create method directly (we now use beta API)
+    with patch.object(client.client.beta.messages, 'create') as mock_create: # type: ignore
         # Create a mock response
         mock_response = Message(
             id="test_id",
@@ -269,10 +269,168 @@ 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_anthropic_client_uses_fast_model_from_provider_manager_without_override():
+    """Test that Anthropic 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 = AnthropicClient(api_key="test_key")
+    
+    # Mock the beta.messages.create method directly (we now use beta API)
+    with patch.object(client.client.beta.messages, 'create') as mock_create: # type: ignore
+        # Create a mock response
+        mock_response = Message(
+            id="test_id",
+            role="assistant",
+            content=[TextBlock(type="text", text="test")], 
+            model='anthropic-model-we-do-not-check', 
+            type="message",
+            usage=Usage(input_tokens=0, output_tokens=0)
+        )
+        mock_create.return_value = mock_response
+
+        # Use a fast model that would be selected by ProviderManager
+        fast_model = get_fast_model_for_selected_model("claude-sonnet-4-5-20250929")
+        
+        await client.request_completions(
+            messages=[{"role": "user", "content": "Test message"}],
+            model=fast_model,
+            message_type=MessageType.CHAT,
+            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_anthropic_client_uses_smartest_model_from_provider_manager_without_override():
+    """Test that Anthropic 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 = AnthropicClient(api_key="test_key")
+    
+    # Mock the beta.messages.create method directly (we now use beta API)
+    with patch.object(client.client.beta.messages, 'create') as mock_create: # type: ignore
+        # Create a mock response
+        mock_response = Message(
+            id="test_id",
+            role="assistant",
+            content=[TextBlock(type="text", text="test")], 
+            model='anthropic-model-we-do-not-check', 
+            type="message",
+            usage=Usage(input_tokens=0, output_tokens=0)
+        )
+        mock_create.return_value = mock_response
+
+        # Use a smartest model that would be selected by ProviderManager
+        smartest_model = get_smartest_model_for_selected_model("claude-haiku-4-5-20251001")
+        
+        await client.request_completions(
+            messages=[{"role": "user", "content": "Test message"}],
+            model=smartest_model,
+            message_type=MessageType.CHAT,
+            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_anthropic_client_stream_uses_fast_model_from_provider_manager_without_override():
+    """Test that Anthropic 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 = AnthropicClient(api_key="test_key")
+    
+    # Mock the beta.messages.create method for streaming
+    with patch.object(client.client.beta.messages, 'create') as mock_create: # type: ignore
+        # Create a mock stream response
+        class MockStreamChunk:
+            def __init__(self, chunk_type, text=""):
+                self.type = chunk_type
+                if chunk_type == "content_block_delta":
+                    self.delta = MagicMock()
+                    self.delta.type = "text_delta"
+                    self.delta.text = text
+        
+        mock_stream = [
+            MockStreamChunk("content_block_delta", "test"),
+            MockStreamChunk("message_stop")
+        ]
+        mock_create.return_value = iter(mock_stream)
+
+        # Use a fast model that would be selected by ProviderManager
+        fast_model = get_fast_model_for_selected_model("claude-sonnet-4-5-20250929")
+        
+        reply_chunks = []
+        def mock_reply(chunk):
+            reply_chunks.append(chunk)
+        
+        await client.stream_completions(
+            messages=[{"role": "user", "content": "Test message"}],
+            model=fast_model,
+            message_id="test-id",
+            message_type=MessageType.CHAT,
+            reply_fn=mock_reply
+        )
+        
+        # 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_anthropic_client_stream_uses_smartest_model_from_provider_manager_without_override():
+    """Test that Anthropic 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 = AnthropicClient(api_key="test_key")
+    
+    # Mock the beta.messages.create method for streaming
+    with patch.object(client.client.beta.messages, 'create') as mock_create: # type: ignore
+        # Create a mock stream response
+        class MockStreamChunk:
+            def __init__(self, chunk_type, text=""):
+                self.type = chunk_type
+                if chunk_type == "content_block_delta":
+                    self.delta = MagicMock()
+                    self.delta.type = "text_delta"
+                    self.delta.text = text
+        
+        mock_stream = [
+            MockStreamChunk("content_block_delta", "test"),
+            MockStreamChunk("message_stop")
+        ]
+        mock_create.return_value = iter(mock_stream)
+
+        # Use a smartest model that would be selected by ProviderManager
+        smartest_model = get_smartest_model_for_selected_model("claude-haiku-4-5-20251001")
+        
+        reply_chunks = []
+        def mock_reply(chunk):
+            reply_chunks.append(chunk)
+        
+        await client.stream_completions(
+            messages=[{"role": "user", "content": "Test message"}],
+            model=smartest_model,
+            message_id="test-id",
+            message_type=MessageType.CHAT,
+            reply_fn=mock_reply
+        )
+        
+        # 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
 
 
 # Caching Tests

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