124 lines
4.4 KiB
Python
124 lines
4.4 KiB
Python
"""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
|