91 lines
3.1 KiB
Python
91 lines
3.1 KiB
Python
"""
|
|
Tests for embedding generation.
|
|
"""
|
|
|
|
import pytest
|
|
import numpy as np
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
|
|
@pytest.mark.integration
|
|
@pytest.mark.embeddings
|
|
class TestOllamaEmbeddings:
|
|
"""Test Ollama embedding generation."""
|
|
|
|
@patch('llama_index.embeddings.ollama.OllamaEmbedding')
|
|
def test_embedding_generation(self, mock_ollama):
|
|
"""Test that embeddings are generated correctly."""
|
|
# Mock the embedding model
|
|
mock_instance = MagicMock()
|
|
mock_instance.get_text_embedding.return_value = np.random.randn(768).tolist()
|
|
mock_ollama.return_value = mock_instance
|
|
|
|
from llama_index.embeddings.ollama import OllamaEmbedding
|
|
|
|
embed_model = OllamaEmbedding(
|
|
model_name="nomic-embed-text",
|
|
base_url="http://localhost:11434"
|
|
)
|
|
|
|
embedding = embed_model.get_text_embedding("test text")
|
|
|
|
assert len(embedding) == 768
|
|
assert all(isinstance(x, float) for x in embedding)
|
|
|
|
def test_embedding_consistency(self):
|
|
"""Test that same text produces same embedding."""
|
|
from llama_index.embeddings.ollama import OllamaEmbedding
|
|
|
|
embed_model = OllamaEmbedding(
|
|
model_name="nomic-embed-text",
|
|
base_url="http://localhost:11434"
|
|
)
|
|
|
|
text = "This is a test sentence for embedding consistency."
|
|
|
|
embedding1 = embed_model.get_text_embedding(text)
|
|
embedding2 = embed_model.get_text_embedding(text)
|
|
|
|
# Should be identical (or very close due to floating point)
|
|
np.testing.assert_array_almost_equal(embedding1, embedding2, decimal=5)
|
|
|
|
def test_embedding_dimension(self):
|
|
"""Test that embeddings have correct dimension."""
|
|
from llama_index.embeddings.ollama import OllamaEmbedding
|
|
|
|
embed_model = OllamaEmbedding(
|
|
model_name="nomic-embed-text",
|
|
base_url="http://localhost:11434"
|
|
)
|
|
|
|
texts = [
|
|
"Short text",
|
|
"This is a longer text with more words and content.",
|
|
"A" * 1000 # Very long text
|
|
]
|
|
|
|
for text in texts:
|
|
embedding = embed_model.get_text_embedding(text)
|
|
assert len(embedding) == 768, f"Expected 768 dimensions, got {len(embedding)}"
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestEmbeddingUtils:
|
|
"""Test embedding utility functions."""
|
|
|
|
def test_cosine_similarity(self):
|
|
"""Test cosine similarity calculation."""
|
|
# Identical vectors should have similarity of 1.0
|
|
vec1 = [1.0, 0.0, 0.0]
|
|
vec2 = [1.0, 0.0, 0.0]
|
|
|
|
similarity = 1 - np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
|
|
assert abs(similarity) < 0.001 # Should be ~0 (distance), so similarity ~1
|
|
|
|
def test_embedding_normalization(self):
|
|
"""Test that embeddings can be normalized."""
|
|
embedding = np.random.randn(768)
|
|
normalized = embedding / np.linalg.norm(embedding)
|
|
|
|
# Normalized vector should have magnitude 1
|
|
assert abs(np.linalg.norm(normalized) - 1.0) < 0.001
|