Livrare LOT 1 - Didi

ai_platform/modules/rerank/tests/__init__.py

@@ -0,0 +1 @@
+"""Tests for rerank module."""

ai_platform/modules/rerank/tests/conftest.py

@@ -0,0 +1,81 @@
+"""Pytest fixtures for rerank tests."""
+
+from unittest.mock import AsyncMock
+
+import pytest
+from fastapi.testclient import TestClient
+
+from rerank.api.app import create_app
+from rerank.api.dependencies import init_concurrency_limiter
+from rerank.client import RerankClient
+from rerank.config import RerankSettings, SettingsCache
+from rerank.schemas import RerankResponse
+from rerank.types import BackendType, RerankResult, RerankUsage
+
+
+@pytest.fixture(autouse=True)
+def clear_settings_cache() -> None:
+    """Clear settings cache before each test."""
+    SettingsCache.clear()
+
+
+@pytest.fixture
+def test_settings() -> RerankSettings:
+    """Create test settings with mocked values."""
+    return RerankSettings(
+        default_backend="vllm",
+        enable_vllm=True,
+        enable_llamacpp=False,
+        external_url="http://localhost:54200",
+        host="127.0.0.1",
+        port=54200,
+        api_tokens=None,
+    )
+
+
+@pytest.fixture
+def mock_rerank_response() -> RerankResponse:
+    """Create a mock rerank response."""
+    return RerankResponse(
+        model="test-model",
+        results=[
+            RerankResult(index=2, relevance_score=0.95),
+            RerankResult(index=0, relevance_score=0.82),
+            RerankResult(index=1, relevance_score=0.10),
+        ],
+        usage=RerankUsage(total_tokens=150),
+        backend="vllm",
+    )
+
+
+@pytest.fixture
+def client_with_mock_backend(
+    test_settings: RerankSettings, mock_rerank_response: RerankResponse
+) -> RerankClient:
+    """Create a RerankClient with mocked backend."""
+    SettingsCache.set(test_settings)
+
+    client = RerankClient(settings=test_settings)
+
+    # Mock the backend's rerank method
+    backend = client.registry.get(BackendType.VLLM)
+    results = [(r.index, r.relevance_score) for r in mock_rerank_response.results]
+    usage = mock_rerank_response.usage
+    backend.rerank = AsyncMock(return_value=(results, usage))
+
+    return client
+
+
+@pytest.fixture
+def app_client(test_settings: RerankSettings) -> TestClient:
+    """Create a FastAPI TestClient with test settings."""
+    SettingsCache.set(test_settings)
+
+    init_concurrency_limiter(test_settings.max_concurrent_reranks)
+
+    app = create_app()
+
+    app.state.settings = test_settings
+    app.state.client = RerankClient(settings=test_settings)
+
+    return TestClient(app)

ai_platform/modules/rerank/tests/test_config.py

@@ -0,0 +1,179 @@
+"""Tests for configuration module."""
+
+import os
+from unittest.mock import patch
+
+import pytest
+from pydantic import ValidationError
+
+from rerank.config import RerankSettings, SettingsCache
+
+
+class TestRerankSettings:
+    """Tests for RerankSettings."""
+
+    def test_required_fields(self) -> None:
+        """Test that required fields must be provided."""
+        with pytest.raises(ValidationError):
+            RerankSettings()
+
+    def test_with_required_fields(self) -> None:
+        """Test settings with all required fields provided."""
+        settings = RerankSettings(
+            default_backend="vllm",
+            enable_vllm=True,
+            enable_llamacpp=False,
+            external_url="http://localhost:54200",
+        )
+
+        assert settings.default_backend == "vllm"
+        assert settings.enable_vllm is True
+        assert settings.enable_llamacpp is False
+        assert settings.external_url == "http://localhost:54200"
+
+    def test_optional_defaults(self) -> None:
+        """Test optional fields have sensible defaults."""
+        settings = RerankSettings(
+            default_backend="vllm",
+            enable_vllm=True,
+            enable_llamacpp=False,
+            external_url="http://localhost:54200",
+        )
+
+        assert settings.host == "0.0.0.0"
+        assert settings.port == 54200
+        assert settings.request_timeout == 120.0
+        assert settings.rate_limit_rps == 20.0
+        assert settings.max_concurrent_reranks == 20
+
+    def test_env_override(self) -> None:
+        """Test environment variable overrides."""
+        env_vars = {
+            "RERANK_DEFAULT_BACKEND": "llamacpp",
+            "RERANK_ENABLE_VLLM": "false",
+            "RERANK_ENABLE_LLAMACPP": "true",
+            "RERANK_EXTERNAL_URL": "http://example.com",
+            "RERANK_PORT": "14200",
+        }
+
+        with patch.dict(os.environ, env_vars, clear=False):
+            settings = RerankSettings()
+
+            assert settings.default_backend == "llamacpp"
+            assert settings.enable_vllm is False
+            assert settings.enable_llamacpp is True
+            assert settings.port == 14200
+
+    def test_vllm_url_default(self) -> None:
+        """Test vLLM URL default value."""
+        settings = RerankSettings(
+            default_backend="vllm",
+            enable_vllm=True,
+            enable_llamacpp=False,
+            external_url="http://localhost:54200",
+        )
+        assert settings.vllm_base_url == "http://localhost:54201"
+
+    def test_llamacpp_url_default(self) -> None:
+        """Test llama.cpp URL default value."""
+        settings = RerankSettings(
+            default_backend="vllm",
+            enable_vllm=True,
+            enable_llamacpp=False,
+            external_url="http://localhost:54200",
+        )
+        assert settings.llamacpp_base_url == "http://localhost:54210"
+
+    def test_api_tokens_parsing(self) -> None:
+        """Test API tokens are parsed from comma-separated string."""
+        settings = RerankSettings(
+            default_backend="vllm",
+            enable_vllm=True,
+            enable_llamacpp=False,
+            external_url="http://localhost:54200",
+            api_tokens="token1,token2,token3",  # type: ignore[arg-type]
+        )
+
+        assert settings.api_tokens is not None
+        assert "token1" in settings.api_tokens
+        assert "token2" in settings.api_tokens
+        assert "token3" in settings.api_tokens
+        assert settings.auth_enabled is True
+
+    def test_auth_disabled_by_default(self) -> None:
+        """Test authentication is disabled when no tokens set."""
+        settings = RerankSettings(
+            default_backend="vllm",
+            enable_vllm=True,
+            enable_llamacpp=False,
+            external_url="http://localhost:54200",
+        )
+
+        assert settings.api_tokens is None
+        assert settings.auth_enabled is False
+
+
+class TestSettingsCache:
+    """Tests for SettingsCache class."""
+
+    def test_set_and_get(self) -> None:
+        """Test setting and getting cached settings."""
+        SettingsCache.clear()
+
+        settings = RerankSettings(
+            default_backend="vllm",
+            enable_vllm=True,
+            enable_llamacpp=False,
+            external_url="http://localhost:54200",
+        )
+        SettingsCache.set(settings)
+
+        retrieved = SettingsCache.get()
+        assert retrieved is settings
+
+    def test_get_without_set_creates_instance(self) -> None:
+        """Test that get() creates instance from environment if not set."""
+        SettingsCache.clear()
+
+        env_vars = {
+            "RERANK_DEFAULT_BACKEND": "vllm",
+            "RERANK_ENABLE_VLLM": "true",
+            "RERANK_ENABLE_LLAMACPP": "false",
+            "RERANK_EXTERNAL_URL": "http://localhost:54200",
+        }
+
+        with patch.dict(os.environ, env_vars, clear=False):
+            settings = SettingsCache.get()
+            assert isinstance(settings, RerankSettings)
+            assert settings.default_backend == "vllm"
+
+    def test_cached_returns_same_instance(self) -> None:
+        """Test that SettingsCache returns cached instance."""
+        SettingsCache.clear()
+
+        settings = RerankSettings(
+            default_backend="vllm",
+            enable_vllm=True,
+            enable_llamacpp=False,
+            external_url="http://localhost:54200",
+        )
+        SettingsCache.set(settings)
+
+        settings1 = SettingsCache.get()
+        settings2 = SettingsCache.get()
+
+        assert settings1 is settings2
+
+    def test_clear_removes_cached_instance(self) -> None:
+        """Test that clear() removes the cached instance."""
+        settings = RerankSettings(
+            default_backend="vllm",
+            enable_vllm=True,
+            enable_llamacpp=False,
+            external_url="http://localhost:54200",
+        )
+        SettingsCache.set(settings)
+
+        SettingsCache.clear()
+
+        assert SettingsCache._instance is None

ai_platform/modules/rerank/tests/test_schemas.py

@@ -0,0 +1,145 @@
+"""Tests for schemas module."""
+
+import pytest
+from pydantic import ValidationError
+
+from rerank.schemas import RerankRequest, RerankResponse
+from rerank.types import BackendType, RerankResult, RerankUsage
+
+
+class TestRerankRequest:
+    """Tests for RerankRequest schema."""
+
+    def test_basic_request(self) -> None:
+        """Test basic request creation."""
+        request = RerankRequest(
+            model="test-model",
+            query="What is AI?",
+            documents=["doc1", "doc2", "doc3"],
+        )
+        assert request.model == "test-model"
+        assert request.query == "What is AI?"
+        assert len(request.documents) == 3
+
+    def test_empty_documents_rejected(self) -> None:
+        """Test that empty documents list is rejected."""
+        with pytest.raises(ValidationError):
+            RerankRequest(
+                model="test-model",
+                query="query",
+                documents=[],
+            )
+
+    def test_empty_string_in_documents_rejected(self) -> None:
+        """Test that empty strings in documents are rejected."""
+        with pytest.raises(ValidationError):
+            RerankRequest(
+                model="test-model",
+                query="query",
+                documents=["doc1", "", "doc3"],
+            )
+
+    def test_empty_query_rejected(self) -> None:
+        """Test that empty query is rejected."""
+        with pytest.raises(ValidationError):
+            RerankRequest(
+                model="test-model",
+                query="",
+                documents=["doc1", "doc2"],
+            )
+
+    def test_top_n_parameter(self) -> None:
+        """Test top_n parameter."""
+        request = RerankRequest(
+            model="test-model",
+            query="query",
+            documents=["doc1", "doc2", "doc3"],
+            top_n=2,
+        )
+        assert request.top_n == 2
+
+    def test_return_documents_parameter(self) -> None:
+        """Test return_documents parameter."""
+        request = RerankRequest(
+            model="test-model",
+            query="query",
+            documents=["doc1", "doc2"],
+            return_documents=True,
+        )
+        assert request.return_documents is True
+
+    def test_backend_override(self) -> None:
+        """Test backend override."""
+        request = RerankRequest(
+            model="test-model",
+            query="query",
+            documents=["doc1", "doc2"],
+            backend=BackendType.LLAMACPP,
+        )
+        assert request.backend == BackendType.LLAMACPP
+
+    def test_invalid_top_n(self) -> None:
+        """Test that invalid top_n is rejected."""
+        with pytest.raises(ValidationError):
+            RerankRequest(
+                model="test-model",
+                query="query",
+                documents=["doc1", "doc2"],
+                top_n=0,
+            )
+
+
+class TestRerankResponse:
+    """Tests for RerankResponse schema."""
+
+    def test_creation(self) -> None:
+        """Test creating response."""
+        response = RerankResponse(
+            model="test-model",
+            results=[
+                RerankResult(index=2, relevance_score=0.95),
+                RerankResult(index=0, relevance_score=0.82),
+            ],
+            usage=RerankUsage(total_tokens=100),
+            backend="vllm",
+        )
+        assert response.model == "test-model"
+        assert len(response.results) == 2
+        assert response.backend == "vllm"
+        assert response.id.startswith("rerank-")
+
+    def test_results_ordering(self) -> None:
+        """Test that results maintain order."""
+        results = [
+            RerankResult(index=2, relevance_score=0.95),
+            RerankResult(index=0, relevance_score=0.82),
+            RerankResult(index=1, relevance_score=0.10),
+        ]
+        response = RerankResponse(
+            model="test-model",
+            results=results,
+            usage=RerankUsage(total_tokens=100),
+            backend="vllm",
+        )
+        assert response.results[0].index == 2
+        assert response.results[0].relevance_score == 0.95
+        assert response.results[1].index == 0
+        assert response.results[2].index == 1
+
+
+class TestRerankResult:
+    """Tests for RerankResult model."""
+
+    def test_without_document(self) -> None:
+        """Test result without document."""
+        result = RerankResult(index=0, relevance_score=0.95)
+        assert result.document is None
+
+    def test_with_document(self) -> None:
+        """Test result with document."""
+        result = RerankResult(
+            index=0,
+            relevance_score=0.95,
+            document="This is the document text",
+        )
+        assert result.document == "This is the document text"