mito-ai 0.1.50__py3-none-any.whl

mito_ai/tests/providers/test_mito_server_utils.py

@@ -0,0 +1,448 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+import json
+import time
+from unittest.mock import MagicMock, patch, AsyncMock
+from tornado.httpclient import HTTPResponse
+
+from mito_ai.utils.mito_server_utils import (
+    ProviderCompletionException,
+    get_response_from_mito_server
+)
+from mito_ai.completions.models import MessageType
+
+
+@pytest.fixture
+def mock_request_params():
+    """Standard request parameters for testing."""
+    return {
+        "url": "https://api.example.com",
+        "headers": {"Content-Type": "application/json"},
+        "data": {"query": "test query"},
+        "timeout": 30,
+        "max_retries": 3,
+        "message_type": MessageType.CHAT,
+        "provider_name": "Test Provider"
+    }
+
+
+@pytest.fixture
+def mock_http_dependencies():
+    """Mock the HTTP client and related dependencies."""
+    with patch('mito_ai.utils.mito_server_utils._create_http_client') as mock_create_client, \
+         patch('mito_ai.utils.mito_server_utils.update_mito_server_quota') as mock_update_quota, \
+         patch('mito_ai.utils.mito_server_utils.check_mito_server_quota') as mock_check_quota, \
+         patch('mito_ai.utils.mito_server_utils.time.time') as mock_time:
+        
+        # Setup mock HTTP client
+        mock_http_client = MagicMock()
+        mock_create_client.return_value = (mock_http_client, 30)
+        
+        # Setup mock time
+        mock_time.side_effect = [0.0, 1.5]  # start_time, end_time
+        
+        yield {
+            'mock_check_quota': mock_check_quota,
+            'mock_create_client': mock_create_client,
+            'mock_http_client': mock_http_client,
+            'mock_update_quota': mock_update_quota,
+            'mock_time': mock_time
+        }
+
+
+def create_mock_response(body_content: dict):
+    """Helper to create mock HTTP response."""
+    mock_response = MagicMock(spec=HTTPResponse)
+    mock_response.body.decode.return_value = json.dumps(body_content)
+    return mock_response
+
+
+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
+
+
+class TestGetResponseFromMitoServer:
+    """Test the get_response_from_mito_server function."""
+    
+    @pytest.mark.parametrize("completion_value,message_type", [
+        ("This is the AI response", MessageType.CHAT),
+        ("Code completion here", MessageType.INLINE_COMPLETION),
+        ("", MessageType.CHAT),  # Empty string
+        (None, MessageType.INLINE_COMPLETION),  # None value
+        ("Multi-line\nresponse\nhere", MessageType.CHAT),  # Multi-line response
+    ])
+    @pytest.mark.asyncio
+    async def test_successful_completion_responses(
+        self, 
+        completion_value,
+        message_type: MessageType,
+        mock_request_params,
+        mock_http_dependencies
+    ):
+        """Test successful responses with various completion values."""
+        # Setup
+        response_body = {"completion": completion_value}
+        mock_response = create_mock_response(response_body)
+        mock_http_dependencies['mock_http_client'].fetch = AsyncMock(return_value=mock_response)
+        
+        # Update request params
+        mock_request_params["message_type"] = message_type
+        
+        # Execute
+        result = await get_response_from_mito_server(**mock_request_params)
+        
+        # Verify
+        assert result == completion_value
+        mock_http_dependencies['mock_check_quota'].assert_called_once_with(message_type)
+        mock_http_dependencies['mock_update_quota'].assert_called_once_with(message_type)
+        mock_http_dependencies['mock_http_client'].close.assert_called_once()
+        
+        # Verify HTTP request was made correctly
+        mock_http_dependencies['mock_http_client'].fetch.assert_called_once_with(
+            mock_request_params["url"],
+            method="POST",
+            headers=mock_request_params["headers"],
+            body=json.dumps(mock_request_params["data"]),
+            request_timeout=30
+        )
+    
+    @pytest.mark.parametrize("error_message,provider_name,expected_exception_provider", [
+        (
+            "There was an error accessing the Anthropic API: Error code: 529 - {'type': 'error', 'error': {'type': 'overloaded_error', 'message': 'Overloaded'}}",
+            "Anthropic",
+            "Anthropic"
+        ),
+        (
+            "Rate limit exceeded",
+            "OpenAI",
+            "OpenAI"
+        ),
+        (
+            "Invalid API key",
+            "Custom Provider",
+            "Custom Provider"
+        ),
+        (
+            "Server timeout",
+            "Mito Server",
+            "Mito Server"
+        ),
+    ])
+    @pytest.mark.asyncio
+    async def test_error_responses_from_server(
+        self, 
+        error_message: str,
+        provider_name: str,
+        expected_exception_provider: str,
+        mock_request_params,
+        mock_http_dependencies
+    ):
+        """Test server returns error response with various error messages and providers."""
+        # Setup
+        response_body = {"error": error_message}
+        mock_response = create_mock_response(response_body)
+        mock_http_dependencies['mock_http_client'].fetch = AsyncMock(return_value=mock_response)
+        
+        # Update request params
+        mock_request_params["provider_name"] = provider_name
+        
+        # Execute and verify exception
+        with pytest.raises(ProviderCompletionException) as exc_info:
+            await get_response_from_mito_server(**mock_request_params)
+        
+        # Verify exception details
+        exception = exc_info.value
+        assert exception.error_message == error_message
+        assert exception.provider_name == expected_exception_provider
+        assert f"{expected_exception_provider} Error" in str(exception)
+        
+        # Verify quota was updated and client was closed
+        mock_http_dependencies['mock_update_quota'].assert_called_once_with(mock_request_params["message_type"])
+        mock_http_dependencies['mock_http_client'].close.assert_called_once()
+    
+    @pytest.mark.parametrize("response_body,expected_error_contains", [
+        ({"some_other_field": "value"}, "No completion found in response"),
+        ({"data": "value", "status": "ok"}, "No completion found in response"),
+        ({}, "No completion found in response"),
+        ({"completion": None, "error": "also present"}, None),  # completion takes precedence
+    ])
+    @pytest.mark.asyncio
+    async def test_invalid_response_formats(
+        self, 
+        response_body: dict,
+        expected_error_contains: str,
+        mock_request_params,
+        mock_http_dependencies
+    ):
+        """Test responses with invalid formats."""
+        # Setup
+        mock_response = create_mock_response(response_body)
+        mock_http_dependencies['mock_http_client'].fetch = AsyncMock(return_value=mock_response)
+        
+        if "completion" in response_body:
+            # This should succeed because completion field exists
+            result = await get_response_from_mito_server(**mock_request_params)
+            assert result == response_body["completion"]
+            mock_http_dependencies['mock_update_quota'].assert_called_once()
+        else:
+            # Execute and verify exception
+            with pytest.raises(ProviderCompletionException) as exc_info:
+                await get_response_from_mito_server(**mock_request_params)
+            
+            # Verify exception details
+            exception = exc_info.value
+            assert expected_error_contains in exception.error_message
+            assert str(response_body) in exception.error_message
+            assert exception.provider_name == mock_request_params["provider_name"]
+            
+            # Verify quota was NOT updated
+            mock_http_dependencies['mock_update_quota'].assert_called_once_with(mock_request_params["message_type"])
+        
+        # Client should always be closed
+        mock_http_dependencies['mock_http_client'].close.assert_called_once()
+    
+    @pytest.mark.parametrize("invalid_json_content,expected_error_contains", [
+        ("invalid json content", "Error parsing response"),
+        ('{"incomplete": json', "Error parsing response"),
+        ("", "Error parsing response"),
+        ('{"malformed":', "Error parsing response"),
+    ])
+    @pytest.mark.asyncio
+    async def test_json_parsing_errors(
+        self, 
+        invalid_json_content: str,
+        expected_error_contains: str,
+        mock_request_params,
+        mock_http_dependencies
+    ):
+        """Test response with invalid or malformed JSON."""
+        # Setup
+        mock_response = MagicMock(spec=HTTPResponse)
+        mock_response.body.decode.return_value = invalid_json_content
+        mock_http_dependencies['mock_http_client'].fetch = AsyncMock(return_value=mock_response)
+        
+        # Execute and verify exception
+        with pytest.raises(ProviderCompletionException) as exc_info:
+            await get_response_from_mito_server(**mock_request_params)
+        
+        # Verify exception details
+        exception = exc_info.value
+        assert expected_error_contains in exception.error_message
+        assert exception.provider_name == mock_request_params["provider_name"]
+        
+        # Verify quota was updated and client was closed
+        mock_http_dependencies['mock_update_quota'].assert_called_once_with(mock_request_params["message_type"])
+        mock_http_dependencies['mock_http_client'].close.assert_called_once()
+    
+    @pytest.mark.parametrize("timeout,max_retries", [
+        (30, 3),
+        (45, 5),
+        (60, 1),
+        (15, 0),
+    ])
+    @pytest.mark.asyncio
+    async def test_http_client_creation_parameters(
+        self, 
+        timeout: int,
+        max_retries: int,
+        mock_request_params,
+        mock_http_dependencies
+    ):
+        """Test that HTTP client is created with correct parameters."""
+        # Setup
+        response_body = {"completion": "test response"}
+        mock_response = create_mock_response(response_body)
+        mock_http_dependencies['mock_http_client'].fetch = AsyncMock(return_value=mock_response)
+        
+        # Update request params
+        mock_request_params["timeout"] = timeout
+        mock_request_params["max_retries"] = max_retries
+        
+        # Execute
+        await get_response_from_mito_server(**mock_request_params)
+        
+        # Verify HTTP client creation
+        mock_http_dependencies['mock_create_client'].assert_called_once_with(timeout, max_retries)
+    
+    @pytest.mark.parametrize("exception_type,exception_message", [
+        (Exception, "Network error"),
+        (ConnectionError, "Connection failed"),
+        (TimeoutError, "Request timed out"),
+        (RuntimeError, "Runtime error occurred"),
+    ])
+    @pytest.mark.asyncio
+    async def test_http_client_always_closed_on_exception(
+        self, 
+        exception_type,
+        exception_message: str,
+        mock_request_params,
+        mock_http_dependencies
+    ):
+        """Test that HTTP client is always closed even when exceptions occur."""
+        # Setup - make fetch raise an exception
+        test_exception = exception_type(exception_message)
+        mock_http_dependencies['mock_http_client'].fetch = AsyncMock(side_effect=test_exception)
+        
+        # Execute and expect exception to bubble up
+        with pytest.raises(exception_type, match=exception_message):
+            await get_response_from_mito_server(**mock_request_params)
+        
+        # Verify client was still closed despite the exception
+        mock_http_dependencies['mock_http_client'].close.assert_called_once()
+    
+    @pytest.mark.asyncio
+    async def test_default_provider_name(self, mock_http_dependencies):
+        """Test that default provider name is used when not specified."""
+        # Setup
+        error_message = "Test error"
+        response_body = {"error": error_message}
+        mock_response = create_mock_response(response_body)
+        mock_http_dependencies['mock_http_client'].fetch = AsyncMock(return_value=mock_response)
+        
+        # Test data without provider_name parameter
+        request_params = {
+            "url": "https://api.example.com",
+            "headers": {"Content-Type": "application/json"},
+            "data": {"query": "test query"},
+            "timeout": 30,
+            "max_retries": 3,
+            "message_type": MessageType.CHAT,
+            # Note: not providing provider_name parameter
+        }
+        
+        # Execute and verify exception
+        with pytest.raises(ProviderCompletionException) as exc_info:
+            await get_response_from_mito_server(**request_params) # type: ignore
+        
+        # Verify default provider name is used
+        exception = exc_info.value
+        assert exception.provider_name == "Mito Server"
+        assert "Mito Server Error" in str(exception)
+    
+    @pytest.mark.asyncio
+    async def test_provider_completion_exception_reraised(self, mock_request_params, mock_http_dependencies):
+        """Test that ProviderCompletionException is re-raised correctly during JSON parsing."""
+        # Setup - simulate ProviderCompletionException during JSON parsing
+        mock_response = MagicMock(spec=HTTPResponse)
+        mock_response.body.decode.return_value = "some json content"  # This will trigger json.loads
+        
+        def mock_json_loads(content, **kwargs):
+            raise ProviderCompletionException("Custom parsing error", "Custom Provider")
+        
+        mock_http_dependencies['mock_http_client'].fetch = AsyncMock(return_value=mock_response)
+        
+        with patch('mito_ai.utils.mito_server_utils.json.loads', side_effect=mock_json_loads), \
+             patch('mito_ai.utils.mito_server_utils.check_mito_server_quota') as mock_check_quota:
+            
+            # Execute and verify exception
+            with pytest.raises(ProviderCompletionException) as exc_info:
+                await get_response_from_mito_server(**mock_request_params)
+            
+            # Verify the original exception is preserved
+            exception = exc_info.value
+            assert exception.error_message == "Custom parsing error"
+            assert exception.provider_name == "Custom Provider"
+            
+            # Verify quota check was called
+            mock_check_quota.assert_called_once_with(mock_request_params["message_type"])
+        
+        # Verify client was closed
+        mock_http_dependencies['mock_http_client'].close.assert_called_once()
+
+    
+    @pytest.mark.parametrize("scenario,response_setup,main_exception,quota_exception", [
+        ("successful_with_quota_error", {"completion": "Success"}, None, Exception("Quota update failed")),
+        ("server_error_with_quota_error", {"error": "Server error"}, ProviderCompletionException, Exception("Quota update failed")),
+        ("invalid_format_with_quota_error", {"invalid": "format"}, ProviderCompletionException, RuntimeError("Quota system down")),
+        ("success_with_quota_timeout", {"completion": "Success"}, None, TimeoutError("Quota service timeout")),
+    ])
+    @pytest.mark.asyncio
+    async def test_quota_update_exceptions_do_not_interfere(
+        self, 
+        scenario: str,
+        response_setup: dict,
+        main_exception,
+        quota_exception,
+        mock_request_params,
+        mock_http_dependencies
+    ):
+        """Test that quota update exceptions don't interfere with main function logic."""
+        # Setup
+        mock_response = create_mock_response(response_setup)
+        mock_http_dependencies['mock_http_client'].fetch = AsyncMock(return_value=mock_response)
+        mock_http_dependencies['mock_update_quota'].side_effect = quota_exception
+        
+        # Execute
+        if main_exception:
+            with pytest.raises(main_exception) as exc_info:
+                await get_response_from_mito_server(**mock_request_params)
+            
+            # Verify the original error is preserved, not the quota error
+            if "error" in response_setup:
+                assert exc_info.value.error_message == response_setup["error"]
+        else:
+            # Should still succeed despite quota update failure
+            result = await get_response_from_mito_server(**mock_request_params)
+            assert result == response_setup["completion"]
+        
+        # Verify quota update was attempted
+        mock_http_dependencies['mock_update_quota'].assert_called_once_with(mock_request_params["message_type"])

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