mito-ai 0.1.33__py3-none-any.whl → 0.1.35__py3-none-any.whl

mito_ai/tests/providers/test_model_resolution.py

@@ -0,0 +1,130 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+"""
+These tests ensure that the correct model is chosen for each message type, for each provider.
+"""
+
+import pytest
+from mito_ai.utils.provider_utils import does_message_require_fast_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 traitlets.config import Config
+
+@pytest.fixture
+def provider_config() -> Config:
+    """Create a proper Config object for the OpenAIProvider."""
+    config = Config()
+    config.OpenAIProvider = Config()
+    config.OpenAIClient = Config()
+    return config
+
+@pytest.fixture
+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
+    
+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
+
+@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."""
+    # 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"
+        
+        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 = 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"
+
+@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."""
+    # 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"
+        
+        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 = 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"

mito_ai/tests/providers/test_openai_client.py

@@ -0,0 +1,57 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+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.asyncio 
+async def test_model_selection_based_on_message_type(message_type, expected_model):
+    """
+    Tests that the correct model is selected based on the message type.
+    """
+    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
+
+        await client.request_completions(
+            message_type=message_type,
+            messages=[{"role": "user", "content": "Test message"}],
+            model=CUSTOM_MODEL,
+            response_format_info=None
+        )
+        
+        # Verify that create was called with the expected model
+        mock_create.assert_called_once()
+        call_args = mock_create.call_args
+        assert call_args[1]['model'] == expected_model

mito_ai/tests/providers/test_provider_completion_exception.py

@@ -0,0 +1,66 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from mito_ai.utils.mito_server_utils import ProviderCompletionException
+import pytest
+
+
+class TestProviderCompletionException:
+    """Test the ProviderCompletionException class."""
+    
+    @pytest.mark.parametrize("error_message,provider_name,error_type,expected_title,expected_hint_contains", [
+        (
+            "Something went wrong", 
+            "LLM Provider", 
+            "LLMProviderError",
+            "LLM Provider Error: Something went wrong",
+            "LLM Provider"
+        ),
+        (
+            "API key is invalid",
+            "OpenAI", 
+            "AuthenticationError",
+            "OpenAI Error: API key is invalid",
+            "OpenAI"
+        ),
+        (
+            "There was an error accessing the Anthropic API: Error code: 529 - {'type': 'error', 'error': {'type': 'overloaded_error', 'message': 'Overloaded'}}",
+            "Anthropic",
+            "LLMProviderError",
+            "Anthropic Error: There was an error accessing the Anthropic API: Error code: 529 - {'type': 'error', 'error': {'type': 'overloaded_error', 'message': 'Overloaded'}}",
+            "Anthropic"
+        ),
+    ])
+    def test_exception_initialization(
+        self, 
+        error_message: str, 
+        provider_name: str, 
+        error_type: str,
+        expected_title: str,
+        expected_hint_contains: str
+    ):
+        """Test exception initialization with various parameter combinations."""
+        exception = ProviderCompletionException(
+            error_message, 
+            provider_name=provider_name, 
+            error_type=error_type
+        )
+        
+        assert exception.error_message == error_message
+        assert exception.provider_name == provider_name
+        assert exception.error_type == error_type
+        assert exception.user_friendly_title == expected_title
+        assert expected_hint_contains in exception.user_friendly_hint
+        assert str(exception) == expected_title
+        assert exception.args[0] == expected_title
+    
+    def test_default_initialization(self):
+        """Test exception initialization with default values."""
+        error_msg = "Something went wrong"
+        exception = ProviderCompletionException(error_msg)
+        
+        assert exception.error_message == error_msg
+        assert exception.provider_name == "LLM Provider"
+        assert exception.error_type == "LLMProviderError"
+        assert exception.user_friendly_title == "LLM Provider Error: Something went wrong"
+        assert "LLM Provider" in exception.user_friendly_hint

mito_ai/tests/providers/test_provider_limits.py

@@ -0,0 +1,42 @@
+# Copyright (c) Saga Inc.
+# 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.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
+
+FAKE_API_KEY = "sk-1234567890"
+
+@pytest.fixture
+def provider_config() -> Config:
+    """Create a proper Config object for the OpenAIProvider."""
+    config = Config()
+    config.OpenAIProvider = Config()
+    config.OpenAIClient = Config()
+    return config
+
+@pytest.mark.parametrize("is_pro,completion_count", [
+    (False, 1),  # OS user below limit
+    (False, OS_MONTHLY_AI_COMPLETIONS_LIMIT + 1),  # OS user above limit
+    (True, 1),   # Pro user below limit
+    (True, OS_MONTHLY_AI_COMPLETIONS_LIMIT + 1),   # Pro user above limit
+])
+def test_openai_provider_with_limits(
+    is_pro: bool, 
+    completion_count: int, 
+    monkeypatch: pytest.MonkeyPatch, 
+    provider_config: Config) -> None:
+    """Test OpenAI provider behavior with different user types and usage limits."""
+    monkeypatch.setenv("OPENAI_API_KEY", FAKE_API_KEY)
+    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", FAKE_API_KEY)
+
+    with (
+        patch_server_limits(is_pro=is_pro, completion_count=completion_count),
+        mock_openai_client()
+    ):
+        llm = OpenAIProvider(config=provider_config)
+        capabilities = llm.capabilities
+        assert "user key" in capabilities.provider
+        assert llm.last_error is None

mito_ai/tests/providers/test_providers.py

@@ -0,0 +1,382 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from __future__ import annotations
+from datetime import datetime
+from typing import Any, List, Optional
+from unittest.mock import patch, MagicMock, AsyncMock
+
+from mito_ai.tests.providers.utils import mock_azure_openai_client, mock_openai_client, patch_server_limits
+import pytest
+from traitlets.config import Config
+from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.completions.models import (
+    MessageType,
+    AICapabilities,
+    CompletionReply
+)
+from mito_ai.utils.server_limits import OS_MONTHLY_AI_COMPLETIONS_LIMIT
+from openai.types.chat import ChatCompletionMessageParam
+
+REALLY_OLD_DATE = "2020-01-01"
+TODAY = datetime.now().strftime("%Y-%m-%d")
+FAKE_API_KEY = "sk-1234567890"
+
+@pytest.fixture
+def provider_config() -> Config:
+    """Create a proper Config object for the OpenAIProvider."""
+    config = Config()
+    config.OpenAIProvider = Config()
+    config.OpenAIClient = Config()
+    return config
+
+@pytest.fixture(autouse=True)
+def reset_env_vars(monkeypatch: pytest.MonkeyPatch) -> None:
+    for var in [
+        "OPENAI_API_KEY", "CLAUDE_API_KEY",
+        "GEMINI_API_KEY", "OLLAMA_MODEL",
+        "AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_MODEL"
+    ]:
+        monkeypatch.delenv(var, raising=False)
+
+
+# ====================
+# TESTS
+# ====================
+
+@pytest.mark.parametrize("provider_config_data", [
+    {
+        "name": "openai",
+        "env_vars": {"OPENAI_API_KEY": FAKE_API_KEY},
+        "constants": {"OPENAI_API_KEY": FAKE_API_KEY},
+        "model": "gpt-4o-mini",
+        "mock_patch": "mito_ai.completions.providers.OpenAIClient",
+        "mock_method": "request_completions",
+        "provider_name": "OpenAI with user key",
+        "key_type": "user"
+    },
+    {
+        "name": "claude", 
+        "env_vars": {"CLAUDE_API_KEY": "claude-key"},
+        "constants": {"CLAUDE_API_KEY": "claude-key", "OPENAI_API_KEY": None},
+        "model": "claude-3-opus-20240229",
+        "mock_patch": "mito_ai.completions.providers.AnthropicClient",
+        "mock_method": "request_completions",
+        "provider_name": "Claude",
+        "key_type": "claude"
+    },
+    {
+        "name": "gemini",
+        "env_vars": {"GEMINI_API_KEY": "gemini-key"},
+        "constants": {"GEMINI_API_KEY": "gemini-key", "OPENAI_API_KEY": None},
+        "model": "gemini-2.0-flash",
+        "mock_patch": "mito_ai.completions.providers.GeminiClient",
+        "mock_method": "request_completions",
+        "provider_name": "Gemini",
+        "key_type": "gemini"
+    },
+    {
+        "name": "azure",
+        "env_vars": {"AZURE_OPENAI_API_KEY": "azure-key"},
+        "constants": {"AZURE_OPENAI_API_KEY": "azure-key", "OPENAI_API_KEY": None},
+        "model": "gpt-4o",
+        "mock_patch": "mito_ai.completions.providers.OpenAIClient",
+        "mock_method": "request_completions",
+        "provider_name": "Azure OpenAI",
+        "key_type": "azure"
+    }
+])
+@pytest.mark.asyncio
+async def test_completion_request(
+    provider_config_data: dict,
+    monkeypatch: pytest.MonkeyPatch, 
+    provider_config: Config
+) -> None:
+    """Test completion requests for different providers."""
+    # Set up environment variables
+    for env_var, value in provider_config_data["env_vars"].items():
+        monkeypatch.setenv(env_var, value)
+    
+    # Set up constants
+    for constant, value in provider_config_data["constants"].items():
+        monkeypatch.setattr(f"mito_ai.constants.{constant}", value)
+
+    # Create mock client
+    mock_client = MagicMock()
+    mock_client.capabilities = AICapabilities(
+        configuration={"model": provider_config_data["model"]},
+        provider=provider_config_data["provider_name"],
+        type="ai_capabilities"
+    )
+    mock_client.key_type = provider_config_data["key_type"]
+    mock_client.request_completions = AsyncMock(return_value="Test completion")
+    mock_client.stream_completions = AsyncMock(return_value="Test completion")
+
+    with patch(provider_config_data["mock_patch"], return_value=mock_client):
+        llm = OpenAIProvider(config=provider_config)
+        messages: List[ChatCompletionMessageParam] = [
+            {"role": "user", "content": "Test message"}
+        ]
+
+        completion = await llm.request_completions(
+            message_type=MessageType.CHAT,
+            messages=messages,
+            model=provider_config_data["model"]
+        )
+
+        assert completion == "Test completion"
+        getattr(mock_client, provider_config_data["mock_method"]).assert_called_once()
+
+
+@pytest.mark.parametrize("provider_config_data", [
+    {
+        "name": "openai",
+        "env_vars": {"OPENAI_API_KEY": FAKE_API_KEY},
+        "constants": {"OPENAI_API_KEY": FAKE_API_KEY},
+        "model": "gpt-4o-mini",
+        "mock_patch": "mito_ai.completions.providers.OpenAIClient",
+        "mock_method": "stream_completions",
+        "provider_name": "OpenAI with user key",
+        "key_type": "user"
+    },
+    {
+        "name": "claude", 
+        "env_vars": {"CLAUDE_API_KEY": "claude-key"},
+        "constants": {"CLAUDE_API_KEY": "claude-key", "OPENAI_API_KEY": None},
+        "model": "claude-3-opus-20240229",
+        "mock_patch": "mito_ai.completions.providers.AnthropicClient",
+        "mock_method": "stream_completions", 
+        "provider_name": "Claude",
+        "key_type": "claude"
+    },
+    {
+        "name": "gemini",
+        "env_vars": {"GEMINI_API_KEY": "gemini-key"},
+        "constants": {"GEMINI_API_KEY": "gemini-key", "OPENAI_API_KEY": None},
+        "model": "gemini-2.0-flash",
+        "mock_patch": "mito_ai.completions.providers.GeminiClient",
+        "mock_method": "stream_completions",
+        "provider_name": "Gemini",
+        "key_type": "gemini"
+    },
+])
+@pytest.mark.asyncio
+async def test_stream_completion_parameterized(
+    provider_config_data: dict,
+    monkeypatch: pytest.MonkeyPatch, 
+    provider_config: Config
+) -> None:
+    """Test stream completions for different providers."""
+    # Set up environment variables
+    for env_var, value in provider_config_data["env_vars"].items():
+        monkeypatch.setenv(env_var, value)
+    
+    # Set up constants
+    for constant, value in provider_config_data["constants"].items():
+        monkeypatch.setattr(f"mito_ai.constants.{constant}", value)
+
+    # Create mock client
+    mock_client = MagicMock()
+    mock_client.capabilities = AICapabilities(
+        configuration={"model": provider_config_data["model"]},
+        provider=provider_config_data["provider_name"],
+        type="ai_capabilities"
+    )
+    mock_client.key_type = provider_config_data["key_type"]
+    mock_client.request_completions = AsyncMock(return_value="Test completion")
+    mock_client.stream_completions = AsyncMock(return_value="Test completion")
+    mock_client.stream_response = AsyncMock(return_value="Test completion")  # For Claude
+
+    with patch(provider_config_data["mock_patch"], return_value=mock_client):
+        llm = OpenAIProvider(config=provider_config)
+        messages: List[ChatCompletionMessageParam] = [
+            {"role": "user", "content": "Test message"}
+        ]
+
+        reply_chunks = []
+        def mock_reply(chunk):
+            reply_chunks.append(chunk)
+
+        completion = await llm.stream_completions(
+            message_type=MessageType.CHAT,
+            messages=messages,
+            model=provider_config_data["model"],
+            message_id="test-id",
+            thread_id="test-thread",
+            reply_fn=mock_reply
+        )
+
+        assert completion == "Test completion"
+        getattr(mock_client, provider_config_data["mock_method"]).assert_called_once()
+        assert len(reply_chunks) > 0
+        assert isinstance(reply_chunks[0], CompletionReply)
+
+
+def test_error_handling(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
+    monkeypatch.setenv("OPENAI_API_KEY", "invalid-key")
+    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", "invalid-key")
+    mock_client = MagicMock()
+    mock_client.capabilities = AICapabilities(
+        configuration={"model": "gpt-4o-mini"},
+        provider="OpenAI with user key",
+        type="ai_capabilities"
+    )
+    mock_client.key_type = "user"
+    mock_client.request_completions.side_effect = Exception("API error")
+
+    with patch("mito_ai.completions.providers.OpenAIClient", return_value=mock_client):
+        llm = OpenAIProvider(config=provider_config)
+        assert llm.last_error is None  # Error should be None until a request is made
+
+def test_claude_error_handling(monkeypatch: pytest.MonkeyPatch, provider_config: Config) -> None:
+    monkeypatch.setenv("CLAUDE_API_KEY", "invalid-key")
+    monkeypatch.setattr("mito_ai.constants.CLAUDE_API_KEY", "invalid-key")
+    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", None)
+
+    mock_client = MagicMock()
+    mock_client.capabilities = AICapabilities(
+        configuration={"model": "claude-3-opus-20240229"},
+        provider="Claude",
+        type="ai_capabilities"
+    )
+    mock_client.key_type = "claude"
+    mock_client.request_completions.side_effect = Exception("API error")
+
+    with patch("mito_ai.completions.providers.AnthropicClient", return_value=mock_client):
+        llm = OpenAIProvider(config=provider_config)
+        assert llm.last_error is None  # Error should be None until a request is made
+
+
+# Mito Server Fallback Tests
+@pytest.mark.parametrize("mito_server_config", [
+    {
+        "name": "openai_fallback",
+        "model": "gpt-4o-mini",
+        "mock_function": "mito_ai.openai_client.get_ai_completion_from_mito_server",
+        "provider_name": "Mito server",
+        "key_type": "mito_server"
+    },
+    {
+        "name": "claude_fallback", 
+        "model": "claude-3-opus-20240229",
+        "mock_function": "mito_ai.anthropic_client.get_anthropic_completion_from_mito_server",
+        "provider_name": "Claude",
+        "key_type": "claude"
+    },
+    {
+        "name": "gemini_fallback",
+        "model": "gemini-2.0-flash",
+        "mock_function": "mito_ai.gemini_client.get_gemini_completion_from_mito_server",
+        "provider_name": "Gemini",
+        "key_type": "gemini"
+    },
+])
+@pytest.mark.asyncio
+async def test_mito_server_fallback_completion_request(
+    mito_server_config: dict,
+    monkeypatch: pytest.MonkeyPatch, 
+    provider_config: Config
+) -> None:
+    """Test that completion requests fallback to Mito server when no API keys are set."""
+    # Clear all API keys to force Mito server fallback
+    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", None)
+    monkeypatch.setattr("mito_ai.constants.CLAUDE_API_KEY", None)
+    monkeypatch.setattr("mito_ai.constants.GEMINI_API_KEY", None)
+    monkeypatch.setattr("mito_ai.enterprise.utils.is_azure_openai_configured", lambda: False)
+    provider_config.OpenAIProvider.api_key = None
+
+    # Mock the appropriate Mito server function
+    with patch(mito_server_config["mock_function"], new_callable=AsyncMock) as mock_mito_function:
+        mock_mito_function.return_value = "Mito server response"
+    
+        messages: List[ChatCompletionMessageParam] = [
+            {"role": "user", "content": "Test message"}
+        ]
+
+        with patch_server_limits():
+                llm = OpenAIProvider(config=provider_config)
+
+                completion = await llm.request_completions(
+                    message_type=MessageType.CHAT,
+                    messages=messages,
+                    model=mito_server_config["model"]
+                )
+
+                assert completion == "Mito server response"
+                mock_mito_function.assert_called_once()
+
+
+@pytest.mark.parametrize("mito_server_config", [
+    {
+        "name": "openai_fallback",
+        "model": "gpt-4o-mini",
+        "mock_function": "mito_ai.openai_client.stream_ai_completion_from_mito_server",
+        "provider_name": "Mito server",
+        "key_type": "mito_server"
+    },
+    {
+        "name": "claude_fallback", 
+        "model": "claude-3-opus-20240229",
+        "mock_function": "mito_ai.anthropic_client.stream_anthropic_completion_from_mito_server",
+        "provider_name": "Claude",
+        "key_type": "claude"
+    },
+    {
+        "name": "gemini_fallback",
+        "model": "gemini-2.0-flash",
+        "mock_function": "mito_ai.gemini_client.stream_gemini_completion_from_mito_server",
+        "provider_name": "Gemini",
+        "key_type": "gemini"
+    },
+])
+@pytest.mark.asyncio
+async def test_mito_server_fallback_stream_completion(
+    mito_server_config: dict,
+    monkeypatch: pytest.MonkeyPatch, 
+    provider_config: Config
+) -> None:
+    """Test that stream completions fallback to Mito server when no API keys are set."""
+    # Clear all API keys to force Mito server fallback
+    monkeypatch.setattr("mito_ai.constants.OPENAI_API_KEY", None)
+    monkeypatch.setattr("mito_ai.constants.CLAUDE_API_KEY", None)
+    monkeypatch.setattr("mito_ai.constants.GEMINI_API_KEY", None)
+    monkeypatch.setattr("mito_ai.enterprise.utils.is_azure_openai_configured", lambda: False)
+    provider_config.OpenAIProvider.api_key = None
+
+    # Create an async generator that yields chunks for streaming
+    async def mock_stream_generator():
+        yield "Chunk 1"
+        yield "Chunk 2"
+        yield "Chunk 3"
+
+    # Mock the appropriate Mito server streaming function
+    with patch(mito_server_config["mock_function"]) as mock_mito_stream:
+        mock_mito_stream.return_value = mock_stream_generator()
+        
+        messages: List[ChatCompletionMessageParam] = [
+            {"role": "user", "content": "Test message"}
+        ]
+        
+        reply_chunks = []
+        def mock_reply(chunk):
+            reply_chunks.append(chunk)
+
+        # Apply patch_server_limits for all cases, not just openai_fallback
+        # Also patch update_mito_server_quota where it's actually used in openai_client
+        with patch_server_limits(), patch("mito_ai.openai_client.update_mito_server_quota", MagicMock(return_value=None)):
+            llm = OpenAIProvider(config=provider_config)                
+
+            completion = await llm.stream_completions(
+                message_type=MessageType.CHAT,
+                messages=messages,
+                model=mito_server_config["model"],
+                message_id="test-id",
+                thread_id="test-thread",
+                reply_fn=mock_reply
+            )
+
+            # Verify that the Mito server function was called
+            mock_mito_stream.assert_called_once()
+            # Verify that reply chunks were generated
+            assert len(reply_chunks) > 0
+            assert isinstance(reply_chunks[0], CompletionReply)