Livrare LOT 1 - Didi

ai_platform/modules/llm-inference/tests/test_client.py

@@ -0,0 +1,201 @@
+"""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