didi-lot1-ai/ai_platform/modules/llm-inference/tests/test_client.py

"""Tests for LLMClient."""

from unittest.mock import AsyncMock, patch

import pytest

from llm_inference.client import LLMClient
from llm_inference.config import LLMSettings, SettingsCache
from llm_inference.exceptions import BackendNotEnabledError
from llm_inference.schemas import CompletionResponse
from llm_inference.types import BackendType, ChatMessage


class TestLLMClient:
    """Tests for LLMClient."""

    def test_init_with_settings(self, test_settings: LLMSettings) -> None:
        """Test client initialization with custom settings."""
        client = LLMClient(settings=test_settings)

        assert client._settings == test_settings
        assert client.registry is not None

    def test_init_without_settings(self, test_settings: LLMSettings) -> None:
        """Test client initialization uses cached settings."""
        # Inject settings into cache
        SettingsCache.set(test_settings)

        client = LLMClient()
        assert client._settings is not None
        assert client._settings.default_backend == test_settings.default_backend

    def test_list_backends(self, test_settings: LLMSettings) -> None:
        """Test listing available backends."""
        client = LLMClient(settings=test_settings)
        backends = client.list_backends()

        # LiteLLM should always be available
        assert BackendType.LITELLM in backends

        # vLLM and llamacpp should not be available (disabled)
        assert BackendType.VLLM not in backends
        assert BackendType.LLAMACPP not in backends

    @pytest.mark.asyncio
    async def test_complete_uses_default_backend(
        self, test_settings: LLMSettings, mock_completion_response: CompletionResponse
    ) -> None:
        """Test that complete uses default backend when not specified."""
        client = LLMClient(settings=test_settings)

        # Mock the backend's complete method
        with patch.object(
            client.registry.get(BackendType.LITELLM),
            "complete",
            new_callable=AsyncMock,
            return_value=mock_completion_response,
        ) as mock_complete:
            response = await client.complete(
                messages=[{"role": "user", "content": "Hello"}],
            )

            assert response.backend == "litellm"
            mock_complete.assert_called_once()

    @pytest.mark.asyncio
    async def test_complete_uses_default_model(
        self, test_settings: LLMSettings, mock_completion_response: CompletionResponse
    ) -> None:
        """Test that complete uses default model when not specified."""
        client = LLMClient(settings=test_settings)

        with patch.object(
            client.registry.get(BackendType.LITELLM),
            "complete",
            new_callable=AsyncMock,
            return_value=mock_completion_response,
        ) as mock_complete:
            await client.complete(
                messages=[{"role": "user", "content": "Hello"}],
            )

            # Check that default model was used
            call_args = mock_complete.call_args
            assert call_args[0][1] == test_settings.default_model

    @pytest.mark.asyncio
    async def test_complete_with_dict_messages(
        self, test_settings: LLMSettings, mock_completion_response: CompletionResponse
    ) -> None:
        """Test that complete accepts dict messages."""
        client = LLMClient(settings=test_settings)

        with patch.object(
            client.registry.get(BackendType.LITELLM),
            "complete",
            new_callable=AsyncMock,
            return_value=mock_completion_response,
        ):
            response = await client.complete(
                messages=[{"role": "user", "content": "Hello"}],
                model="gpt-4",
            )

            assert response is not None

    @pytest.mark.asyncio
    async def test_complete_with_chatmessage_objects(
        self, test_settings: LLMSettings, mock_completion_response: CompletionResponse
    ) -> None:
        """Test that complete accepts ChatMessage objects."""
        client = LLMClient(settings=test_settings)

        with patch.object(
            client.registry.get(BackendType.LITELLM),
            "complete",
            new_callable=AsyncMock,
            return_value=mock_completion_response,
        ):
            response = await client.complete(
                messages=[ChatMessage(role="user", content="Hello")],
                model="gpt-4",
            )

            assert response is not None

    @pytest.mark.asyncio
    async def test_list_models(self, test_settings: LLMSettings) -> None:
        """Test listing models from all backends."""
        # Set API keys so models are returned
        test_settings.openai_api_key = "test-key"
        client = LLMClient(settings=test_settings)

        models = await client.list_models()

        # Should return some models
        assert len(models) > 0
        # All should be from litellm backend
        assert all(m.backend == "litellm" for m in models)

    @pytest.mark.asyncio
    async def test_health_check(self, test_settings: LLMSettings) -> None:
        """Test health check returns backend status."""
        client = LLMClient(settings=test_settings)

        health = await client.health_check()

        # Should return status for litellm backend
        assert "litellm" in health
        # Health check may return True or False depending on connectivity
        assert isinstance(health["litellm"], bool)

    def test_parse_messages_dict(self, test_settings: LLMSettings) -> None:
        """Test message parsing from dicts."""
        client = LLMClient(settings=test_settings)

        messages = client._parse_messages([{"role": "user", "content": "Hello"}])

        assert len(messages) == 1
        assert isinstance(messages[0], ChatMessage)
        assert messages[0].role == "user"
        assert messages[0].content == "Hello"

    def test_parse_messages_chatmessage(self, test_settings: LLMSettings) -> None:
        """Test message parsing from ChatMessage objects."""
        client = LLMClient(settings=test_settings)
        original = ChatMessage(role="user", content="Hello")

        messages = client._parse_messages([original])

        assert len(messages) == 1
        assert messages[0] is original


class TestBackendSelection:
    """Tests for backend selection."""

    def test_unavailable_backend_raises_error(self, test_settings: LLMSettings) -> None:
        """Test that requesting unavailable backend raises error."""
        client = LLMClient(settings=test_settings)

        with pytest.raises(BackendNotEnabledError):
            client.registry.get(BackendType.VLLM)

    def test_enabled_vllm_backend(self) -> None:
        """Test that vLLM backend is available when enabled."""
        settings = LLMSettings(
            default_backend="litellm",
            enable_vllm=True,
            enable_llamacpp=False,
            external_url="http://localhost:8100",
        )

        # This will fail because openai package may not be installed
        # but we can at least verify the registry tries to load it
        client = LLMClient(settings=settings)

        # vLLM might be in backends if openai package is available
        backends = client.list_backends()
        # Just verify litellm is there
        assert BackendType.LITELLM in backends