mito-ai 0.1.50__py3-none-any.whl

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']}"

mito_ai/tests/providers/test_gemini_client.py

@@ -0,0 +1,195 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+import ast
+import inspect
+import requests
+from mito_ai.gemini_client import GeminiClient, get_gemini_system_prompt_and_messages
+from mito_ai.utils.gemini_utils import get_gemini_completion_function_params, FAST_GEMINI_MODEL
+from google.genai.types import Part, GenerateContentResponse, Candidate, Content
+from mito_ai.completions.models import ResponseFormatInfo, AgentResponse
+from unittest.mock import MagicMock, patch
+from typing import List, Dict, Any
+from mito_ai.completions.models import MessageType
+
+# Dummy base64 image (1x1 PNG)
+DUMMY_IMAGE_DATA_URL = (
+    "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/wcAAgMBAp9l9AAAAABJRU5ErkJggg=="
+)
+
+def test_mixed_text_and_image():
+    messages: List[Dict[str, Any]] = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": [
+            {"type": "text", "text": "Here is an image:"},
+            {"type": "image_url", "image_url": {"url": DUMMY_IMAGE_DATA_URL}}
+        ]}
+    ]
+    system_instructions, contents = get_gemini_system_prompt_and_messages(messages)
+    assert system_instructions == "You are a helpful assistant."
+    assert len(contents) == 1
+    assert contents[0]["role"] == "user"
+    # Should have two parts: text and image
+    assert len(contents[0]["parts"]) == 2
+    assert contents[0]["parts"][0]["text"] == "Here is an image:"
+    # The second part should be a Part object (from google.genai.types)
+    assert isinstance(contents[0]["parts"][1], Part)
+
+def test_no_system_instructions_only_content():
+    messages: List[Dict[str, Any]] = [
+        {"role": "user", "content": "Hello!"},
+        {"role": "assistant", "content": "Hi, how can I help you?"}
+    ]
+    system_instructions, contents = get_gemini_system_prompt_and_messages(messages)
+    assert system_instructions == ""
+    assert len(contents) == 2
+    assert contents[0]["role"] == "user"
+    assert contents[0]["parts"][0]["text"] == "Hello!"
+    assert contents[1]["role"] == "model"
+    assert contents[1]["parts"][0]["text"] == "Hi, how can I help you?"
+
+def test_system_instructions_and_content():
+    messages: List[Dict[str, Any]] = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "What is the weather today?"}
+    ]
+    system_instructions, contents = get_gemini_system_prompt_and_messages(messages)
+    assert system_instructions == "You are a helpful assistant."
+    assert len(contents) == 1
+    assert contents[0]["role"] == "user"
+    assert contents[0]["parts"][0]["text"] == "What is the weather today?"
+
+@pytest.mark.asyncio
+async def test_json_response_handling():
+    # Create a mock response with JSON content
+    mock_response = GenerateContentResponse(
+        candidates=[
+            Candidate(
+                content=Content(
+                    parts=[Part(text='{"key": "value"}')]
+                )
+            )
+        ]
+    )
+
+    # Create a mock client with the response
+    client = GeminiClient(api_key="test_key")
+    client.client = MagicMock()
+    client.client.models.generate_content.return_value = mock_response
+
+    # Test with response format info
+    response_format_info = ResponseFormatInfo(name="agent_response", format=AgentResponse)
+    result = await client.request_completions(
+        messages=[{"role": "user", "content": "Test message"}],
+        model="test-model",
+        response_format_info=response_format_info
+    )
+    assert result == '{"key": "value"}'
+
+@pytest.mark.asyncio
+async def test_json_response_handling_with_invalid_json():
+    """
+    Tests how the GeminiClient handles responses with invalid JSON when a response format is specified.
+    
+    This test is important because:
+    1. It verifies that the client doesn't crash when receiving malformed JSON responses
+    2. It ensures that the raw response text is returned even when JSON parsing would fail
+    3. It tests the error handling path in the response processing logic    
+    """
+    # Create a mock response with invalid JSON content
+    mock_response = GenerateContentResponse(
+        candidates=[
+            Candidate(
+                content=Content(
+                    parts=[Part(text='{"key": value}')]  # Invalid JSON - missing quotes around 'value'
+                )
+            )
+        ]
+    )
+
+    # Create a mock client with the response
+    client = GeminiClient(api_key="test_key")
+    client.client = MagicMock()
+    client.client.models.generate_content.return_value = mock_response
+
+    # Test with response format info
+    response_format_info = ResponseFormatInfo(name="agent_response", format=AgentResponse)
+    result = await client.request_completions(
+        messages=[{"role": "user", "content": "Test message"}],
+        model="test-model",
+        response_format_info=response_format_info
+    )
+    # Should return the raw string even if JSON is invalid
+    assert result == '{"key": value}'
+
+@pytest.mark.asyncio
+async def test_json_response_handling_with_multiple_parts():
+    # Create a mock response with multiple parts
+    mock_response = GenerateContentResponse(
+        candidates=[
+            Candidate(
+                content=Content(
+                    parts=[
+                        Part(text='Here is the JSON: '),
+                        Part(text='{"key": "value"}'),
+                        Part(text=' End of response')
+                    ]
+                )
+            )
+        ]
+    )
+
+    # Create a mock client with the response
+    client = GeminiClient(api_key="test_key")
+    client.client = MagicMock()
+    client.client.models.generate_content.return_value = mock_response
+
+    # Test with response format info
+    response_format_info = ResponseFormatInfo(name="agent_response", format=AgentResponse)
+    result = await client.request_completions(
+        messages=[{"role": "user", "content": "Test message"}],
+        model="test-model",
+        response_format_info=response_format_info
+    )
+    # Should concatenate all parts
+    assert result == 'Here is the JSON: {"key": "value"} End of response' 
+    
+CUSTOM_MODEL = "smart-gemini-model"
+@pytest.mark.parametrize("message_type, expected_model", [
+    (MessageType.CHAT, CUSTOM_MODEL),  #
+    (MessageType.SMART_DEBUG, CUSTOM_MODEL),  #
+    (MessageType.CODE_EXPLAIN, CUSTOM_MODEL),  #
+    (MessageType.AGENT_EXECUTION, CUSTOM_MODEL),  #
+    (MessageType.AGENT_AUTO_ERROR_FIXUP, CUSTOM_MODEL),  #
+    (MessageType.INLINE_COMPLETION, FAST_GEMINI_MODEL),  #
+    (MessageType.CHAT_NAME_GENERATION, FAST_GEMINI_MODEL),  #
+])
+@pytest.mark.asyncio 
+async def test_get_completion_model_selection_based_on_message_type(message_type, expected_model):
+    """
+    Tests that the correct model is selected based on the message type.
+    """
+    with patch('google.genai.Client') as mock_genai_class:
+        mock_client = MagicMock()
+        mock_models = MagicMock()
+        mock_client.models = mock_models
+        mock_genai_class.return_value = mock_client
+        
+        client = GeminiClient(api_key="test_key")
+        
+        # Create a mock response
+        mock_response = 'test-response'
+        mock_models.generate_content.return_value = mock_response
+
+        await client.request_completions(
+            messages=[{"role": "user", "content": "Test message"}],
+            model=CUSTOM_MODEL,
+            message_type=message_type,
+            response_format_info=None
+        )
+        
+        # Verify that generate_content was called with the expected model
+        mock_models.generate_content.assert_called_once()
+        call_args = mock_models.generate_content.call_args
+        assert call_args[1]['model'] == expected_model