Livrare LOT 1 - Didi

ai_platform/modules/embeddings/tests/test_schemas.py

@@ -0,0 +1,124 @@
+"""Tests for schemas module."""
+
+import pytest
+from pydantic import ValidationError
+
+from embeddings.schemas import (
+    EmbeddingRequest,
+    EmbeddingResponse,
+    encode_embedding_base64,
+)
+from embeddings.types import BackendType, EmbeddingData, EmbeddingUsage
+
+
+class TestEmbeddingRequest:
+    """Tests for EmbeddingRequest schema."""
+
+    def test_single_string_input(self) -> None:
+        """Test that single string input is converted to list."""
+        request = EmbeddingRequest(input="hello", model="test-model")
+        assert request.input == ["hello"]
+
+    def test_list_input(self) -> None:
+        """Test list input is preserved."""
+        request = EmbeddingRequest(input=["hello", "world"], model="test-model")
+        assert request.input == ["hello", "world"]
+
+    def test_empty_input_rejected(self) -> None:
+        """Test that empty input is rejected."""
+        with pytest.raises(ValidationError):
+            EmbeddingRequest(input=[], model="test-model")
+
+    def test_empty_string_in_list_rejected(self) -> None:
+        """Test that empty strings in list are rejected."""
+        with pytest.raises(ValidationError):
+            EmbeddingRequest(input=["hello", ""], model="test-model")
+
+    def test_default_encoding_format(self) -> None:
+        """Test default encoding format is float."""
+        request = EmbeddingRequest(input="hello", model="test-model")
+        assert request.encoding_format == "float"
+
+    def test_base64_encoding_format(self) -> None:
+        """Test base64 encoding format."""
+        request = EmbeddingRequest(
+            input="hello", model="test-model", encoding_format="base64"
+        )
+        assert request.encoding_format == "base64"
+
+    def test_backend_override(self) -> None:
+        """Test backend override."""
+        request = EmbeddingRequest(
+            input="hello", model="test-model", backend=BackendType.LLAMACPP
+        )
+        assert request.backend == BackendType.LLAMACPP
+
+    def test_dimensions_parameter(self) -> None:
+        """Test dimensions parameter."""
+        request = EmbeddingRequest(input="hello", model="test-model", dimensions=256)
+        assert request.dimensions == 256
+
+    def test_invalid_dimensions(self) -> None:
+        """Test that invalid dimensions are rejected."""
+        with pytest.raises(ValidationError):
+            EmbeddingRequest(input="hello", model="test-model", dimensions=0)
+
+
+class TestEmbeddingResponse:
+    """Tests for EmbeddingResponse schema."""
+
+    def test_creation(self) -> None:
+        """Test creating response."""
+        response = EmbeddingResponse(
+            data=[EmbeddingData(index=0, embedding=[0.1, 0.2, 0.3])],
+            model="test-model",
+            usage=EmbeddingUsage(prompt_tokens=5, total_tokens=5),
+            backend="vllm",
+        )
+        assert response.object == "list"
+        assert len(response.data) == 1
+        assert response.model == "test-model"
+        assert response.backend == "vllm"
+
+    def test_multiple_embeddings(self) -> None:
+        """Test response with multiple embeddings."""
+        response = EmbeddingResponse(
+            data=[
+                EmbeddingData(index=0, embedding=[0.1, 0.2]),
+                EmbeddingData(index=1, embedding=[0.3, 0.4]),
+            ],
+            model="test-model",
+            usage=EmbeddingUsage(prompt_tokens=10, total_tokens=10),
+            backend="vllm",
+        )
+        assert len(response.data) == 2
+        assert response.data[0].index == 0
+        assert response.data[1].index == 1
+
+
+class TestEncodeEmbeddingBase64:
+    """Tests for base64 encoding function."""
+
+    def test_encode_simple(self) -> None:
+        """Test encoding simple embedding."""
+        embedding = [1.0, 2.0, 3.0]
+        encoded = encode_embedding_base64(embedding)
+        assert isinstance(encoded, str)
+        # Should be base64 encoded
+        assert len(encoded) > 0
+
+    def test_encode_decode_roundtrip(self) -> None:
+        """Test that encoding can be reversed."""
+        import base64
+        import struct
+
+        embedding = [0.1, 0.2, 0.3, 0.4, 0.5]
+        encoded = encode_embedding_base64(embedding)
+
+        # Decode
+        decoded_bytes = base64.b64decode(encoded)
+        decoded = list(struct.unpack(f"<{len(embedding)}f", decoded_bytes))
+
+        # Compare with tolerance for float precision
+        for original, decoded_val in zip(embedding, decoded):
+            assert abs(original - decoded_val) < 1e-6