mito-ai 0.1.32__py3-none-any.whl → 0.1.34__py3-none-any.whl

mito_ai/gemini_client.py

@@ -5,10 +5,9 @@ from typing import Any, Callable, Dict, List, Optional, Union, Tuple
 from google import genai
 from google.genai import types
 from google.genai.types import GenerateContentConfig, Part, Content, GenerateContentResponse
-from mito_ai.completions.models import CompletionItem, CompletionReply, CompletionStreamChunk, MessageType, ResponseFormatInfo
+from mito_ai.completions.models import CompletionError, CompletionItem, CompletionReply, CompletionStreamChunk, MessageType, ResponseFormatInfo
 from mito_ai.utils.gemini_utils import get_gemini_completion_from_mito_server, stream_gemini_completion_from_mito_server, get_gemini_completion_function_params
-
-GEMINI_FAST_MODEL = "gemini-2.0-flash-lite"
+from mito_ai.utils.mito_server_utils import ProviderCompletionException
 
 def extract_and_parse_gemini_json_response(response: GenerateContentResponse) -> Optional[str]:
     """
@@ -100,65 +99,62 @@ def get_gemini_system_prompt_and_messages(messages: List[Dict[str, Any]]) -> Tup
 
 
 class GeminiClient:
-    def __init__(self, api_key: Optional[str], model: str):
+    def __init__(self, api_key: Optional[str]):
         self.api_key = api_key
-        self.model = model
         if api_key:
             self.client = genai.Client(api_key=api_key)
 
     async def request_completions(
         self,
         messages: List[Dict[str, Any]],
+        model: str,
         response_format_info: Optional[ResponseFormatInfo] = None,
         message_type: MessageType = MessageType.CHAT
     ) -> str:
-        try:
-            # Extract system instructions and contents
-            system_instructions, contents = get_gemini_system_prompt_and_messages(messages)
-
-            # Get provider data for Gemini completion
-            provider_data = get_gemini_completion_function_params(
-                model=self.model if response_format_info else GEMINI_FAST_MODEL,
-                contents=contents,
+        # Extract system instructions and contents
+        system_instructions, contents = get_gemini_system_prompt_and_messages(messages)
+
+        # Get provider data for Gemini completion
+        provider_data = get_gemini_completion_function_params(
+            model=model,
+            contents=contents,
+            message_type=message_type,
+            response_format_info=response_format_info
+        )
+
+        if self.api_key:
+            # Generate content using the Gemini client
+            response_config = GenerateContentConfig(
+                system_instruction=system_instructions,
+                response_mime_type=provider_data.get("config", {}).get("response_mime_type"),
+                response_schema=provider_data.get("config", {}).get("response_schema")
+            )
+            response = self.client.models.generate_content(
+                model=provider_data["model"],
+                contents=contents,  # type: ignore
+                config=response_config
+            )
+            
+            result = extract_and_parse_gemini_json_response(response)
+            
+            if not result:
+                return "No response received from Gemini API"
+
+            return result
+        else:
+            # Fallback to Mito server for completion
+            return await get_gemini_completion_from_mito_server(
+                model=provider_data["model"],
+                contents=messages, # Use the extracted contents instead of converted messages to avoid serialization issues
                 message_type=message_type,
-                response_format_info=response_format_info
+                config=provider_data.get("config", None),
+                response_format_info=response_format_info,
             )
 
-            if self.api_key:
-                # Generate content using the Gemini client
-                response_config = GenerateContentConfig(
-                    system_instruction=system_instructions,
-                    response_mime_type=provider_data.get("config", {}).get("response_mime_type"),
-                    response_schema=provider_data.get("config", {}).get("response_schema")
-                )
-                response = self.client.models.generate_content(
-                    model=provider_data["model"],
-                    contents=contents,  # type: ignore
-                    config=response_config
-                )
-                
-                result = extract_and_parse_gemini_json_response(response)
-                
-                if not result:
-                    return "No response received from Gemini API"
-
-                return result
-            else:
-                # Fallback to Mito server for completion
-                return await get_gemini_completion_from_mito_server(
-                    model=provider_data["model"],
-                    contents=messages, # Use the extracted contents instead of converted messages to avoid serialization issues
-                    message_type=message_type,
-                    config=provider_data.get("config", None),
-                    response_format_info=response_format_info,
-                )
-
-        except Exception as e:
-            return f"Error generating content: {str(e)}"
-
     async def stream_completions(
             self,
             messages: List[Dict[str, Any]],
+            model: str,
             message_id: str,
             reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
             message_type: MessageType = MessageType.CHAT
@@ -169,7 +165,7 @@ class GeminiClient:
             system_instructions, contents = get_gemini_system_prompt_and_messages(messages)
             if self.api_key:
                 for chunk in self.client.models.generate_content_stream(
-                        model=self.model,
+                        model=model,
                         contents=contents,  # type: ignore
                         config=GenerateContentConfig(
                             system_instruction=system_instructions
@@ -208,7 +204,7 @@ class GeminiClient:
                 return accumulated_response
             else:
                 async for chunk_text in stream_gemini_completion_from_mito_server(
-                        model=self.model,
+                        model=model,
                         contents=messages,  # Use the extracted contents instead of converted messages to avoid serialization issues
                         message_type=message_type,
                         message_id=message_id,

mito_ai/openai_client.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Union
 
+from mito_ai.utils.mito_server_utils import ProviderCompletionException
 import openai
 from openai.types.chat import ChatCompletionMessageParam
 from traitlets import Instance, Unicode, default, validate
@@ -36,8 +37,6 @@ from mito_ai.utils.telemetry_utils import (
 
 OPENAI_MODEL_FALLBACK = "gpt-4.1"
 
-OPENAI_FAST_MODEL = "gpt-4.1-nano"
-
 class OpenAIClient(LoggingConfigurable):
     """Provide AI feature through OpenAI services."""
 
@@ -222,26 +221,20 @@ This attribute is observed by the websocket provider to push the error to the cl
         )
         return client
 
-    def _resolve_model(self, model: Optional[str] = None, response_format_info: Optional[ResponseFormatInfo] = None) -> str:
+    def _adjust_model_for_azure_or_ollama(self, model: str) -> str:
         
         # If they have set an Azure OpenAI model, then we always use it
         if is_azure_openai_configured() and constants.AZURE_OPENAI_MODEL is not None:
             self.log.debug(f"Resolving to Azure OpenAI model: {constants.AZURE_OPENAI_MODEL}")
             return constants.AZURE_OPENAI_MODEL
         
-        # Otherwise, we use the fast model for anything other than the agent mode
-        if response_format_info:
-            return OPENAI_FAST_MODEL
-        
         # If they have set an Ollama model, then we use it
         if constants.OLLAMA_MODEL is not None:
             return constants.OLLAMA_MODEL
         
-        # If they have set a model, then we use it
-        if model:
-            return model
+        # Otherwise, we use the model they provided
+        return model
         
-        return OPENAI_MODEL_FALLBACK
 
     async def request_completions(
             self,
@@ -263,39 +256,33 @@ This attribute is observed by the websocket provider to push the error to the cl
         # Reset the last error
         self.last_error = None
         completion = None
+        
+        # Note: We don't catch exceptions here because we want them to bubble up 
+        # to the providers file so we can handle all client exceptions in one place.
 
-        try:
-            
-            # Make sure we are using the correct model
-            # TODO: If we bring back inline completions or another action that needs to 
-            # respond fast, we must require the user to configure a fast model with Azure as well. 
-            model = self._resolve_model(model, response_format_info)
-
-            # Handle other providers as before
-            completion_function_params = get_open_ai_completion_function_params(
-                model, messages, False, response_format_info
+        # Handle other providers as before
+        completion_function_params = get_open_ai_completion_function_params(
+            message_type, model, messages, False, response_format_info
+        )
+        
+        # If they have set an Azure OpenAI or Ollama model, then we use it
+        completion_function_params["model"] = self._adjust_model_for_azure_or_ollama(completion_function_params["model"])
+
+        if self._active_async_client is not None:
+            response = await self._active_async_client.chat.completions.create(**completion_function_params)
+            completion = response.choices[0].message.content or ""
+        else: 
+            last_message_content = str(messages[-1].get("content", "")) if messages else None
+            completion = await get_ai_completion_from_mito_server(
+                last_message_content,
+                completion_function_params,
+                self.timeout,
+                self.max_retries,
+                message_type,
             )
 
-            if self._active_async_client is not None:
-                response = await self._active_async_client.chat.completions.create(**completion_function_params)
-                completion = response.choices[0].message.content or ""
-            else: 
-                last_message_content = str(messages[-1].get("content", "")) if messages else None
-                completion = await get_ai_completion_from_mito_server(
-                    last_message_content,
-                    completion_function_params,
-                    self.timeout,
-                    self.max_retries,
-                    message_type,
-                )
-
-                update_mito_server_quota(message_type)
+        return completion
 
-            return completion
-                
-        except BaseException as e:
-            self.last_error = CompletionError.from_exception(e)
-            raise
     
     async def stream_completions(
         self,
@@ -315,9 +302,6 @@ This attribute is observed by the websocket provider to push the error to the cl
         # Reset the last error
         self.last_error = None
         accumulated_response = ""
-        
-        # Validate that the model is supported.
-        model = self._resolve_model(model, response_format_info)
             
         # Send initial acknowledgment
         reply_fn(CompletionReply(
@@ -329,8 +313,10 @@ This attribute is observed by the websocket provider to push the error to the cl
 
         # Handle other providers as before
         completion_function_params = get_open_ai_completion_function_params(
-            model, messages, True, response_format_info
+            message_type, model, messages, True, response_format_info
         )
+        
+        completion_function_params["model"] = self._adjust_model_for_azure_or_ollama(completion_function_params["model"])
 
         try:
             if self._active_async_client is not None:

mito_ai/tests/message_history/test_generate_short_chat_name.py

@@ -6,10 +6,6 @@ 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.completions.models import MessageType
-from mito_ai.openai_client import OPENAI_FAST_MODEL
-from mito_ai.anthropic_client import ANTHROPIC_FAST_MODEL  
-from mito_ai.gemini_client import GEMINI_FAST_MODEL
 
 
 @pytest.fixture

mito_ai/tests/open_ai_utils_test.py

@@ -80,28 +80,24 @@ def test_prepare_request_data_and_headers_basic() -> None:
         mock_get_user_field.side_effect = ["test@example.com", "user123"]
         
         # Mock the quota check
-        with patch("mito_ai.utils.open_ai_utils.check_mito_server_quota") as mock_check_quota:
-            data, headers = _prepare_request_data_and_headers(
-                last_message_content="test message",
-                ai_completion_data={"key": "value"},
-                timeout=30,
-                max_retries=3,
-                message_type=MessageType.CHAT
-            )
-
-            # Verify quota check was called
-            mock_check_quota.assert_called_once_with(MessageType.CHAT)
-
-            # Verify data structure
-            assert data["timeout"] == 30
-            assert data["max_retries"] == 3
-            assert data["email"] == "test@example.com"
-            assert data["user_id"] == "user123"
-            assert data["data"] == {"key": "value"}
-            assert data["user_input"] == "test message"
-
-            # Verify headers
-            assert headers == {"Content-Type": "application/json"}
+        data, headers = _prepare_request_data_and_headers(
+            last_message_content="test message",
+            ai_completion_data={"key": "value"},
+            timeout=30,
+            max_retries=3,
+            message_type=MessageType.CHAT
+        )
+
+        # Verify data structure
+        assert data["timeout"] == 30
+        assert data["max_retries"] == 3
+        assert data["email"] == "test@example.com"
+        assert data["user_id"] == "user123"
+        assert data["data"] == {"key": "value"}
+        assert data["user_input"] == "test message"
+
+        # Verify headers
+        assert headers == {"Content-Type": "application/json"}
 
 def test_prepare_request_data_and_headers_null_message() -> None:
     """Test handling of null message content"""

mito_ai/tests/{test_anthropic_client.py → providers/test_anthropic_client.py}

@@ -2,15 +2,16 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 import pytest
-from mito_ai.anthropic_client import get_anthropic_system_prompt_and_messages, extract_and_parse_anthropic_json_response, AnthropicClient, ANTHROPIC_FAST_MODEL
-from mito_ai.utils.anthropic_utils import get_anthropic_completion_function_params
-from anthropic.types import MessageParam, Message, ContentBlock, TextBlock, ToolUseBlock, Usage
+from mito_ai.anthropic_client import get_anthropic_system_prompt_and_messages, extract_and_parse_anthropic_json_response, AnthropicClient
+from mito_ai.utils.anthropic_utils import get_anthropic_completion_function_params, 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 ResponseFormatInfo, AgentResponse
+from mito_ai.completions.models import MessageType, ResponseFormatInfo, AgentResponse
 from unittest.mock import MagicMock, patch
 import anthropic
 from typing import List, Dict, Any, cast, Union
 
+
 # Dummy base64 image (1x1 PNG)
 DUMMY_IMAGE_DATA_URL = (
     "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/wcAAgMBAp9l9AAAAABJRU5ErkJggg=="
@@ -231,40 +232,44 @@ def test_tool_use_without_agent_response():
         extract_and_parse_anthropic_json_response(response)
     assert "No valid AgentResponse format found" in str(exc_info.value)
 
-CUSTOM_MODEL = "claude-3-5-sonnet-latest"
-@pytest.mark.parametrize("response_format_info, expected_model", [
-    (ResponseFormatInfo(name="agent_response", format=AgentResponse), CUSTOM_MODEL),  # With response_format_info - should use self.model
-    (None, ANTHROPIC_FAST_MODEL),  # Without response_format_info - should use ANTHROPIC_FAST_MODEL
+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.asyncio
-async def test_model_selection_based_on_response_format_info(response_format_info, expected_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 whether response_format_info is provided.
+    Tests that the correct model is selected based on the message type.
     """
+    client = AnthropicClient(api_key="test_key")
     
-    # Create an AnthropicClient with a specific model
-    custom_model = CUSTOM_MODEL
-    client = AnthropicClient(api_key="test_key", model=custom_model)
-    
-    # Mock the messages.create method to avoid actual API calls
-    client.client = MagicMock()
-    mock_response = Message(
-        id="test_id",
-        role="assistant",
-        content=[TextBlock(type="text", text="Test response")],
-        model=custom_model,
-        type="message",
-        usage=Usage(input_tokens=0, output_tokens=0)
-    )
-    client.client.messages.create.return_value = mock_response
-    
-    with patch('mito_ai.anthropic_client.get_anthropic_completion_function_params', wraps=get_anthropic_completion_function_params) as mock_get_params:
+    # Mock the messages.create method directly
+    with patch.object(client.client.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
+
         await client.request_completions(
             messages=[{"role": "user", "content": "Test message"}],
-            response_format_info=response_format_info
+            model=CUSTOM_MODEL,
+            message_type=message_type,
+            response_format_info=None
         )
         
-        # Verify that get_anthropic_completion_function_params was called with the expected model
-        mock_get_params.assert_called_once()
-        call_args = mock_get_params.call_args
+        # 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_azure.py

@@ -176,15 +176,11 @@ class TestAzureOpenAIClientCreation:
             openai_client = OpenAIClient(config=provider_config)
             
             # Test with gpt-4.1 model
-            resolved_model = openai_client._resolve_model("gpt-4.1")
+            resolved_model = openai_client._adjust_model_for_azure_or_ollama("gpt-4.1")
             assert resolved_model == FAKE_AZURE_MODEL
             
             # Test with any other model
-            resolved_model = openai_client._resolve_model("gpt-3.5-turbo")
-            assert resolved_model == FAKE_AZURE_MODEL
-            
-            # Test with no model specified
-            resolved_model = openai_client._resolve_model()
+            resolved_model = openai_client._adjust_model_for_azure_or_ollama("gpt-3.5-turbo")
             assert resolved_model == FAKE_AZURE_MODEL
 
 

mito_ai/tests/providers/test_capabilities.py

@@ -0,0 +1,120 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+from unittest.mock import MagicMock, patch
+from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.tests.providers.utils import mock_azure_openai_client, mock_openai_client, patch_server_limits
+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("test_case", [
+    {
+        "name": "mito_server_fallback_no_keys",
+        "setup": {
+            "OPENAI_API_KEY": None,
+            "CLAUDE_API_KEY": None, 
+            "GEMINI_API_KEY": None,
+            "is_azure_configured": False,
+        },
+        "expected_provider": "Mito server",
+        "expected_key_type": "mito_server_key"
+    },
+    {
+        "name": "claude_when_only_claude_key",
+        "setup": {
+            "OPENAI_API_KEY": None,
+            "CLAUDE_API_KEY": "claude-test-key",
+            "GEMINI_API_KEY": None,
+            "is_azure_configured": False,
+        },
+        "expected_provider": "Claude",
+        "expected_key_type": "claude"
+    },
+    {
+        "name": "gemini_when_only_gemini_key",
+        "setup": {
+            "OPENAI_API_KEY": None,
+            "CLAUDE_API_KEY": None,
+            "GEMINI_API_KEY": "gemini-test-key",
+            "is_azure_configured": False,
+        },
+        "expected_provider": "Gemini", 
+        "expected_key_type": "gemini"
+    },
+    {
+        "name": "openai_when_openai_key",
+        "setup": {
+            "OPENAI_API_KEY": 'openai-test-key',
+            "CLAUDE_API_KEY": None,
+            "GEMINI_API_KEY": None,
+            "is_azure_configured": False,
+        },
+        "expected_provider": "OpenAI (user key)",
+        "expected_key_type": "user_key"
+    },
+    {
+        "name": "claude_priority_over_gemini",
+        "setup": {
+            "OPENAI_API_KEY": None,
+            "CLAUDE_API_KEY": "claude-test-key",
+            "GEMINI_API_KEY": "gemini-test-key",
+            "is_azure_configured": False,
+        },
+        "expected_provider": "Claude",
+        "expected_key_type": "claude"
+    },
+])
+def test_provider_capabilities_real_logic(
+    test_case: dict,
+    monkeypatch: pytest.MonkeyPatch, 
+    provider_config: Config
+) -> None:
+    """Test the actual provider selection logic in OpenAIProvider.capabilities"""
+    
+    # Set up the environment based on test case
+    setup = test_case["setup"]
+    
+    # CRITICAL: Set up ALL mocks BEFORE creating any clients
+    for key, value in setup.items():
+        if key == "is_azure_configured":
+            if value:
+                # For Azure case, mock to return True and set required constants
+                monkeypatch.setattr("mito_ai.enterprise.utils.is_azure_openai_configured", lambda: True)
+                monkeypatch.setattr("mito_ai.constants.AZURE_OPENAI_MODEL", "gpt-4o")
+            else:
+                # For non-Azure case, mock to return False
+                monkeypatch.setattr("mito_ai.enterprise.utils.is_azure_openai_configured", lambda: False)
+        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:
+            with patch("openai.AsyncAzureOpenAI") as mock_async_azure_openai:
+                # Mock successful API key validation for OpenAI
+                mock_openai_instance = MagicMock()
+                mock_openai_instance.models.list.return_value = [MagicMock(id="gpt-4o-mini")]
+                mock_openai_constructor.return_value = mock_openai_instance
+                
+                # 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)
+                    
+                    # Test capabilities 
+                    capabilities = llm.capabilities
+                    assert capabilities.provider == test_case["expected_provider"], f"Test case: {test_case['name']}"
+                    assert llm.key_type == test_case["expected_key_type"], f"Test case: {test_case['name']}"