diff --git a/api/tests/unit_tests/core/rag/embedding/test_embedding_service.py b/api/tests/unit_tests/core/rag/embedding/test_embedding_service.py index 168d0e466d4..2c395ea9807 100644 --- a/api/tests/unit_tests/core/rag/embedding/test_embedding_service.py +++ b/api/tests/unit_tests/core/rag/embedding/test_embedding_service.py @@ -1,58 +1,24 @@ -"""Comprehensive unit tests for embedding service (CacheEmbedding). +"""Behavior-focused tests for :class:`CacheEmbedding`. -This test module covers all aspects of the embedding service including: -- Batch embedding generation with proper batching logic -- Embedding model switching and configuration -- Embedding dimension validation -- Error handling for API failures -- Cache management (database and Redis) -- Normalization and NaN handling - -Test Coverage: -============== -1. **Batch Embedding Generation** - - Single text embedding - - Multiple texts in batches - - Large batch processing (respects MAX_CHUNKS) - - Empty text handling - -2. **Embedding Model Switching** - - Different providers (OpenAI, Cohere, etc.) - - Different models within same provider - - Model instance configuration - -3. **Embedding Dimension Validation** - - Correct dimensions for different models - - Vector normalization - - Dimension consistency across batches - -4. **Error Handling** - - API connection failures - - Rate limit errors - - Authorization errors - - Invalid input handling - - NaN value detection and handling - -5. **Cache Management** - - Database cache for document embeddings - - Redis cache for query embeddings - - Cache hit/miss scenarios - - Cache invalidation - -All tests use mocking to avoid external dependencies and ensure fast, reliable execution. -Tests follow the Arrange-Act-Assert pattern for clarity. +Document and multimodal caches are persisted in SQLite so cache hits, misses, +provider isolation, uniqueness races, and rollback behavior exercise real +SQLAlchemy statements. Model providers and Redis remain mocked external I/O. """ import base64 -import logging +import pickle +from dataclasses import dataclass from decimal import Decimal from unittest.mock import Mock, patch import numpy as np import pytest +from sqlalchemy import event, select from sqlalchemy.exc import IntegrityError +from sqlalchemy.orm import Session from core.entities.embedding_type import EmbeddingInputType +from core.rag.embedding import cached_embedding from core.rag.embedding.cached_embedding import CacheEmbedding from graphon.model_runtime.entities.model_entities import ModelPropertyKey from graphon.model_runtime.entities.text_embedding_entities import EmbeddingResult, EmbeddingUsage @@ -61,1833 +27,281 @@ from graphon.model_runtime.errors.invoke import ( InvokeConnectionError, InvokeRateLimitError, ) +from libs import helper from models.dataset import Embedding -class TestCacheEmbeddingDocuments: - """Test suite for CacheEmbedding.embed_documents method. +@dataclass(frozen=True) +class _Database: + """Expose a real session through the interface used by the cache.""" - This class tests the batch embedding generation functionality including: - - Single and multiple text processing - - Cache hit/miss scenarios - - Batch processing with MAX_CHUNKS - - Database cache management - - Error handling during embedding generation - """ + session: Session - @pytest.fixture - def mock_model_instance(self): - """Create a mock ModelInstance for testing. - Returns: - Mock: Configured ModelInstance with text embedding capabilities - """ - model_instance = Mock() - model_instance.model_name = "text-embedding-ada-002" - model_instance.provider = "openai" - model_instance.credentials = {"api_key": "test-key"} +@pytest.fixture +def sqlite_embedding_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: + """Bind the shared SQLite session to the document-cache implementation.""" - # Mock the model type instance - model_type_instance = Mock() - model_instance.model_type_instance = model_type_instance + monkeypatch.setattr(cached_embedding, "db", _Database(sqlite_session)) + return sqlite_session - # Mock model schema with MAX_CHUNKS property - model_schema = Mock() - model_schema.model_properties = {ModelPropertyKey.MAX_CHUNKS: 10} - model_type_instance.get_model_schema.return_value = model_schema - return model_instance +@pytest.fixture +def model_instance() -> Mock: + instance = Mock() + instance.model_name = "text-embedding-ada-002" + instance.provider = "openai" + instance.credentials = {"api_key": "test-key"} + schema = Mock(model_properties={ModelPropertyKey.MAX_CHUNKS: 10}) + instance.model_type_instance.get_model_schema.return_value = schema + return instance - @pytest.fixture - def sample_embedding_result(self): - """Create a sample EmbeddingResult for testing. - Returns: - EmbeddingResult: Mock embedding result with proper structure - """ - # Create normalized embedding vectors (dimension 1536 for ada-002) - embedding_vector = np.random.randn(1536) - normalized_vector = (embedding_vector / np.linalg.norm(embedding_vector)).tolist() +def _vector(dimension: int = 8, *, offset: float = 0) -> np.ndarray: + return np.arange(1 + offset, dimension + 1 + offset, dtype=float) - usage = EmbeddingUsage( - tokens=10, - total_tokens=10, + +def _result(*vectors: np.ndarray) -> EmbeddingResult: + return EmbeddingResult( + model="text-embedding-ada-002", + embeddings=list(vectors), + usage=EmbeddingUsage( + tokens=len(vectors), + total_tokens=len(vectors), unit_price=Decimal("0.0001"), price_unit=Decimal(1000), total_price=Decimal("0.000001"), currency="USD", - latency=0.5, - ) - - return EmbeddingResult( - model="text-embedding-ada-002", - embeddings=[normalized_vector], - usage=usage, - ) - - def test_embed_single_document_cache_miss(self, mock_model_instance, sample_embedding_result): - """Test embedding a single document when cache is empty. - - Verifies: - - Model invocation with correct parameters - - Embedding normalization - - Database cache storage - - Correct return value - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = ["Python is a programming language"] - - # Mock database query to return no cached embedding (cache miss) - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - - # Mock model invocation - mock_model_instance.invoke_text_embedding.return_value = sample_embedding_result - - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - assert len(result) == 1 - assert isinstance(result[0], list) - assert len(result[0]) == 1536 # ada-002 dimension - assert all(isinstance(x, float) for x in result[0]) - - # Verify model was invoked with correct parameters - mock_model_instance.invoke_text_embedding.assert_called_once_with( - texts=texts, - input_type=EmbeddingInputType.DOCUMENT, - ) - - # Verify embedding was added to database cache - mock_session.add.assert_called_once() - mock_session.commit.assert_called_once() - - def test_embed_multiple_documents_cache_miss(self, mock_model_instance): - """Test embedding multiple documents when cache is empty. - - Verifies: - - Batch processing of multiple texts - - Multiple embeddings returned - - All embeddings are properly normalized - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = [ - "Python is a programming language", - "JavaScript is used for web development", - "Machine learning is a subset of AI", - ] - - # Create multiple embedding vectors - embeddings = [] - for _ in range(3): - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - embeddings.append(normalized) - - usage = EmbeddingUsage( - tokens=30, - total_tokens=30, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.000003"), - currency="USD", - latency=0.8, - ) - - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=embeddings, - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result - - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - assert len(result) == 3 - assert all(len(emb) == 1536 for emb in result) - assert all(isinstance(emb, list) for emb in result) - - # Verify all embeddings are normalized (L2 norm ≈ 1.0) - for emb in result: - norm = np.linalg.norm(emb) - assert abs(norm - 1.0) < 0.01 # Allow small floating point error - - def test_embed_documents_cache_hit(self, mock_model_instance): - """Test embedding documents when embeddings are already cached. - - Verifies: - - Cached embeddings are retrieved from database - - Model is not invoked for cached texts - - Correct embeddings are returned - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = ["Python is a programming language"] - - # Create cached embedding - cached_vector = np.random.randn(1536) - normalized_cached = (cached_vector / np.linalg.norm(cached_vector)).tolist() - - mock_cached_embedding = Mock(spec=Embedding) - mock_cached_embedding.get_embedding.return_value = normalized_cached - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - # Mock database to return cached embedding (cache hit) - mock_session.scalar.return_value = mock_cached_embedding - - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - assert len(result) == 1 - assert result[0] == normalized_cached - - # Verify model was NOT invoked (cache hit) - mock_model_instance.invoke_text_embedding.assert_not_called() - - # Verify no new cache entries were added - mock_session.add.assert_not_called() - - def test_embed_documents_partial_cache_hit(self, mock_model_instance): - """Test embedding documents with mixed cache hits and misses. - - Verifies: - - Cached embeddings are used when available - - Only non-cached texts are sent to model - - Results are properly merged - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = [ - "Cached text 1", - "New text 1", - "New text 2", - ] - - # Create cached embedding for first text - cached_vector = np.random.randn(1536) - normalized_cached = (cached_vector / np.linalg.norm(cached_vector)).tolist() - - mock_cached_embedding = Mock(spec=Embedding) - mock_cached_embedding.get_embedding.return_value = normalized_cached - - # Create new embeddings for non-cached texts - new_embeddings = [] - for _ in range(2): - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - new_embeddings.append(normalized) - - usage = EmbeddingUsage( - tokens=20, - total_tokens=20, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.000002"), - currency="USD", - latency=0.6, - ) - - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=new_embeddings, - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - with patch("core.rag.embedding.cached_embedding.helper.generate_text_hash") as mock_hash: - # Mock hash generation to return predictable values - hash_counter = [0] - - def generate_hash(text): - hash_counter[0] += 1 - return f"hash_{hash_counter[0]}" - - mock_hash.side_effect = generate_hash - - # Mock database to return cached embedding only for first text (hash_1) - mock_session.scalar.side_effect = [mock_cached_embedding, None, None] - mock_model_instance.invoke_text_embedding.return_value = embedding_result - - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - assert len(result) == 3 - assert result[0] == normalized_cached # From cache - # The model returns already normalized embeddings, but the code normalizes again - # So we just verify the structure and dimensions - assert result[1] is not None - assert isinstance(result[1], list) - assert len(result[1]) == 1536 - assert result[2] is not None - assert isinstance(result[2], list) - assert len(result[2]) == 1536 - - # Verify all embeddings are normalized - for emb in result: - if emb is not None: - norm = np.linalg.norm(emb) - assert abs(norm - 1.0) < 0.01 - - # Verify model was invoked only for non-cached texts - mock_model_instance.invoke_text_embedding.assert_called_once() - call_args = mock_model_instance.invoke_text_embedding.call_args - assert len(call_args.kwargs["texts"]) == 2 # Only 2 non-cached texts - - def test_embed_documents_large_batch(self, mock_model_instance): - """Test embedding a large batch of documents respecting MAX_CHUNKS. - - Verifies: - - Large batches are split according to MAX_CHUNKS - - Multiple model invocations for large batches - - All embeddings are returned correctly - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - # Create 25 texts, MAX_CHUNKS is 10, so should be 3 batches (10, 10, 5) - texts = [f"Text number {i}" for i in range(25)] - - # Create embeddings for each batch - def create_batch_result(batch_size): - embeddings = [] - for _ in range(batch_size): - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - embeddings.append(normalized) - - usage = EmbeddingUsage( - tokens=batch_size * 10, - total_tokens=batch_size * 10, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal(str(batch_size * 0.000001)), - currency="USD", - latency=0.5, - ) - - return EmbeddingResult( - model="text-embedding-ada-002", - embeddings=embeddings, - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - - # Mock model to return appropriate batch results - batch_results = [ - create_batch_result(10), - create_batch_result(10), - create_batch_result(5), - ] - mock_model_instance.invoke_text_embedding.side_effect = batch_results - - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - assert len(result) == 25 - assert all(len(emb) == 1536 for emb in result) - - # Verify model was invoked 3 times (for 3 batches) - assert mock_model_instance.invoke_text_embedding.call_count == 3 - - # Verify batch sizes - calls = mock_model_instance.invoke_text_embedding.call_args_list - assert len(calls[0].kwargs["texts"]) == 10 - assert len(calls[1].kwargs["texts"]) == 10 - assert len(calls[2].kwargs["texts"]) == 5 - - def test_embed_documents_nan_handling(self, mock_model_instance, caplog: pytest.LogCaptureFixture): - """Test handling of NaN values in embeddings. - - Verifies: - - NaN values are detected - - NaN embeddings are skipped - - Warning is logged - - Valid embeddings are still processed - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = ["Valid text", "Text that produces NaN"] - - # Create one valid embedding and one with NaN - # Note: The code normalizes again, so we provide unnormalized vector - valid_vector = np.random.randn(1536) - - # Create NaN vector - nan_vector = [float("nan")] * 1536 - - usage = EmbeddingUsage( - tokens=20, - total_tokens=20, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.000002"), - currency="USD", - latency=0.5, - ) - - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=[valid_vector.tolist(), nan_vector], - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result - - with caplog.at_level(logging.WARNING, logger="core.rag.embedding.cached_embedding"): - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - # NaN embedding is skipped, so only 1 embedding in result - # The first position gets the valid embedding, second is None - assert len(result) == 2 - assert result[0] is not None - assert isinstance(result[0], list) - assert len(result[0]) == 1536 - # Second embedding should be None since NaN was skipped - assert result[1] is None - - # Verify warning was logged - assert sum(1 for r in caplog.records if r.levelno == logging.WARNING) >= 1 - assert any("Normalized embedding is nan" in record.message for record in caplog.records) - - def test_embed_documents_api_connection_error(self, mock_model_instance): - """Test handling of API connection errors during embedding. - - Verifies: - - Connection errors are propagated - - Database transaction is rolled back - - Error message is preserved - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = ["Test text"] - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - - # Mock model to raise connection error - mock_model_instance.invoke_text_embedding.side_effect = InvokeConnectionError("Failed to connect to API") - - # Act & Assert - with pytest.raises(InvokeConnectionError) as exc_info: - cache_embedding.embed_documents(texts) - - assert "Failed to connect to API" in str(exc_info.value) - - # Verify database rollback was called - mock_session.rollback.assert_called() - - def test_embed_documents_rate_limit_error(self, mock_model_instance): - """Test handling of rate limit errors during embedding. - - Verifies: - - Rate limit errors are propagated - - Database transaction is rolled back - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = ["Test text"] - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - - # Mock model to raise rate limit error - mock_model_instance.invoke_text_embedding.side_effect = InvokeRateLimitError("Rate limit exceeded") - - # Act & Assert - with pytest.raises(InvokeRateLimitError) as exc_info: - cache_embedding.embed_documents(texts) - - assert "Rate limit exceeded" in str(exc_info.value) - mock_session.rollback.assert_called() - - def test_embed_documents_authorization_error(self, mock_model_instance): - """Test handling of authorization errors during embedding. - - Verifies: - - Authorization errors are propagated - - Database transaction is rolled back - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = ["Test text"] - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - - # Mock model to raise authorization error - mock_model_instance.invoke_text_embedding.side_effect = InvokeAuthorizationError("Invalid API key") - - # Act & Assert - with pytest.raises(InvokeAuthorizationError) as exc_info: - cache_embedding.embed_documents(texts) - - assert "Invalid API key" in str(exc_info.value) - mock_session.rollback.assert_called() - - def test_embed_documents_database_integrity_error(self, mock_model_instance, sample_embedding_result): - """Test handling of database integrity errors during cache storage. - - Verifies: - - Integrity errors are caught (e.g., duplicate hash) - - Database transaction is rolled back - - Embeddings are still returned - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = ["Test text"] - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_text_embedding.return_value = sample_embedding_result - - # Mock database commit to raise IntegrityError - mock_session.commit.side_effect = IntegrityError("Duplicate key", None, None) - - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - # Embeddings should still be returned despite cache error - assert len(result) == 1 - assert isinstance(result[0], list) - - # Verify rollback was called - mock_session.rollback.assert_called() - - -class TestCacheEmbeddingQuery: - """Test suite for CacheEmbedding.embed_query method. - - This class tests the query embedding functionality including: - - Single query embedding - - Redis cache management - - Cache hit/miss scenarios - - Error handling - """ - - @pytest.fixture - def mock_model_instance(self): - """Create a mock ModelInstance for testing.""" - model_instance = Mock() - model_instance.model_name = "text-embedding-ada-002" - model_instance.provider = "openai" - model_instance.credentials = {"api_key": "test-key"} - return model_instance - - def test_embed_query_cache_miss(self, mock_model_instance): - """Test embedding a query when Redis cache is empty. - - Verifies: - - Model invocation with QUERY input type - - Embedding normalization - - Redis cache storage - - Correct return value - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - query = "What is Python?" - - # Create embedding result - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - - usage = EmbeddingUsage( - tokens=5, - total_tokens=5, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.0000005"), - currency="USD", - latency=0.3, - ) - - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=[normalized], - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: - # Mock Redis cache miss - mock_redis.get.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result - - # Act - result = cache_embedding.embed_query(query) - - # Assert - assert isinstance(result, list) - assert len(result) == 1536 - assert all(isinstance(x, float) for x in result) - - # Verify model was invoked with QUERY input type - mock_model_instance.invoke_text_embedding.assert_called_once_with( - texts=[query], - input_type=EmbeddingInputType.QUERY, - ) - - # Verify Redis cache was set - mock_redis.setex.assert_called_once() - # Cache key format: {provider}_{model}_{hash} - cache_key = mock_redis.setex.call_args[0][0] - assert "openai" in cache_key - assert "text-embedding-ada-002" in cache_key - - # Verify cache TTL is 600 seconds - assert mock_redis.setex.call_args[0][1] == 600 - - def test_embed_query_cache_hit(self, mock_model_instance): - """Test embedding a query when Redis cache contains the result. - - Verifies: - - Cached embedding is retrieved from Redis - - Model is not invoked - - Cache TTL is extended - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - query = "What is Python?" - - # Create cached embedding - vector = np.random.randn(1536) - normalized = vector / np.linalg.norm(vector) - - # Encode to base64 (as stored in Redis) - vector_bytes = normalized.tobytes() - encoded_vector = base64.b64encode(vector_bytes).decode("utf-8") - - with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: - # Mock Redis cache hit - mock_redis.get.return_value = encoded_vector - - # Act - result = cache_embedding.embed_query(query) - - # Assert - assert isinstance(result, list) - assert len(result) == 1536 - - # Verify model was NOT invoked (cache hit) - mock_model_instance.invoke_text_embedding.assert_not_called() - - # Verify cache TTL was extended - mock_redis.expire.assert_called_once() - assert mock_redis.expire.call_args[0][1] == 600 - - def test_embed_query_nan_handling(self, mock_model_instance): - """Test handling of NaN values in query embeddings. - - Verifies: - - NaN values are detected - - ValueError is raised - - Error message is descriptive - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - query = "Query that produces NaN" - - # Create NaN embedding - nan_vector = [float("nan")] * 1536 - - usage = EmbeddingUsage( - tokens=5, - total_tokens=5, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.0000005"), - currency="USD", - latency=0.3, - ) - - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=[nan_vector], - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: - mock_redis.get.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result - - # Act & Assert - with pytest.raises(ValueError) as exc_info: - cache_embedding.embed_query(query) - - assert "Normalized embedding is nan" in str(exc_info.value) - - def test_embed_query_connection_error(self, mock_model_instance): - """Test handling of connection errors during query embedding. - - Verifies: - - Connection errors are propagated - - Error is logged in debug mode - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - query = "Test query" - - with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: - mock_redis.get.return_value = None - - # Mock model to raise connection error - mock_model_instance.invoke_text_embedding.side_effect = InvokeConnectionError("Connection failed") - - # Act & Assert - with pytest.raises(InvokeConnectionError) as exc_info: - cache_embedding.embed_query(query) - - assert "Connection failed" in str(exc_info.value) - - def test_embed_query_redis_cache_error(self, mock_model_instance): - """Test handling of Redis cache errors during storage. - - Verifies: - - Redis errors are caught - - Embedding is still returned - - Error is logged in debug mode - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - query = "Test query" - - # Create valid embedding - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - - usage = EmbeddingUsage( - tokens=5, - total_tokens=5, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.0000005"), - currency="USD", - latency=0.3, - ) - - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=[normalized], - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: - mock_redis.get.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result - - # Mock Redis setex to raise error - mock_redis.setex.side_effect = Exception("Redis connection failed") - - # Act & Assert - with pytest.raises(Exception) as exc_info: - cache_embedding.embed_query(query) - - assert "Redis connection failed" in str(exc_info.value) - - -class TestEmbeddingModelSwitching: - """Test suite for embedding model switching functionality. - - This class tests the ability to switch between different embedding models - and providers, ensuring proper configuration and dimension handling. - """ - - def test_switch_between_openai_models(self): - """Test switching between different OpenAI embedding models. - - Verifies: - - Different models produce different cache keys - - Model name is correctly used in cache lookup - - Embeddings are model-specific - """ - # Arrange - model_instance_ada = Mock() - model_instance_ada.model_name = "text-embedding-ada-002" - model_instance_ada.provider = "openai" - - # Mock model type instance for ada - model_type_instance_ada = Mock() - model_instance_ada.model_type_instance = model_type_instance_ada - model_schema_ada = Mock() - model_schema_ada.model_properties = {ModelPropertyKey.MAX_CHUNKS: 10} - model_type_instance_ada.get_model_schema.return_value = model_schema_ada - - model_instance_3_small = Mock() - model_instance_3_small.model_name = "text-embedding-3-small" - model_instance_3_small.provider = "openai" - - # Mock model type instance for 3-small - model_type_instance_3_small = Mock() - model_instance_3_small.model_type_instance = model_type_instance_3_small - model_schema_3_small = Mock() - model_schema_3_small.model_properties = {ModelPropertyKey.MAX_CHUNKS: 10} - model_type_instance_3_small.get_model_schema.return_value = model_schema_3_small - - cache_ada = CacheEmbedding(model_instance_ada) - cache_3_small = CacheEmbedding(model_instance_3_small) - - text = "Test text" - - # Create different embeddings for each model - vector_ada = np.random.randn(1536) - normalized_ada = (vector_ada / np.linalg.norm(vector_ada)).tolist() - - vector_3_small = np.random.randn(1536) - normalized_3_small = (vector_3_small / np.linalg.norm(vector_3_small)).tolist() - - usage = EmbeddingUsage( - tokens=5, - total_tokens=5, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.0000005"), - currency="USD", - latency=0.3, - ) - - result_ada = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=[normalized_ada], - usage=usage, - ) - - result_3_small = EmbeddingResult( - model="text-embedding-3-small", - embeddings=[normalized_3_small], - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - - model_instance_ada.invoke_text_embedding.return_value = result_ada - model_instance_3_small.invoke_text_embedding.return_value = result_3_small - - # Act - embedding_ada = cache_ada.embed_documents([text]) - embedding_3_small = cache_3_small.embed_documents([text]) - - # Assert - # Both should return embeddings but they should be different - assert len(embedding_ada) == 1 - assert len(embedding_3_small) == 1 - assert embedding_ada[0] != embedding_3_small[0] - - # Verify both models were invoked - model_instance_ada.invoke_text_embedding.assert_called_once() - model_instance_3_small.invoke_text_embedding.assert_called_once() - - def test_switch_between_providers(self): - """Test switching between different embedding providers. - - Verifies: - - Different providers use separate cache namespaces - - Provider name is correctly used in cache lookup - """ - # Arrange - model_instance_openai = Mock() - model_instance_openai.model_name = "text-embedding-ada-002" - model_instance_openai.provider = "openai" - - model_instance_cohere = Mock() - model_instance_cohere.model_name = "embed-english-v3.0" - model_instance_cohere.provider = "cohere" - - cache_openai = CacheEmbedding(model_instance_openai) - cache_cohere = CacheEmbedding(model_instance_cohere) - - query = "Test query" - - # Create embeddings - vector_openai = np.random.randn(1536) - normalized_openai = (vector_openai / np.linalg.norm(vector_openai)).tolist() - - vector_cohere = np.random.randn(1024) # Cohere uses different dimension - normalized_cohere = (vector_cohere / np.linalg.norm(vector_cohere)).tolist() - - usage_openai = EmbeddingUsage( - tokens=5, - total_tokens=5, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.0000005"), - currency="USD", - latency=0.3, - ) - - usage_cohere = EmbeddingUsage( - tokens=5, - total_tokens=5, - unit_price=Decimal("0.0002"), - price_unit=Decimal(1000), - total_price=Decimal("0.000001"), - currency="USD", - latency=0.4, - ) - - result_openai = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=[normalized_openai], - usage=usage_openai, - ) - - result_cohere = EmbeddingResult( - model="embed-english-v3.0", - embeddings=[normalized_cohere], - usage=usage_cohere, - ) - - with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: - mock_redis.get.return_value = None - - model_instance_openai.invoke_text_embedding.return_value = result_openai - model_instance_cohere.invoke_text_embedding.return_value = result_cohere - - # Act - embedding_openai = cache_openai.embed_query(query) - embedding_cohere = cache_cohere.embed_query(query) - - # Assert - assert len(embedding_openai) == 1536 # OpenAI dimension - assert len(embedding_cohere) == 1024 # Cohere dimension - - # Verify different cache keys were used - calls = mock_redis.setex.call_args_list - assert len(calls) == 2 - cache_key_openai = calls[0][0][0] - cache_key_cohere = calls[1][0][0] - - assert "openai" in cache_key_openai - assert "cohere" in cache_key_cohere - assert cache_key_openai != cache_key_cohere - - -class TestEmbeddingDimensionValidation: - """Test suite for embedding dimension validation. - - This class tests that embeddings maintain correct dimensions - and are properly normalized across different scenarios. - """ - - @pytest.fixture - def mock_model_instance(self): - """Create a mock ModelInstance for testing.""" - model_instance = Mock() - model_instance.model_name = "text-embedding-ada-002" - model_instance.provider = "openai" - model_instance.credentials = {"api_key": "test-key"} - - model_type_instance = Mock() - model_instance.model_type_instance = model_type_instance - - model_schema = Mock() - model_schema.model_properties = {ModelPropertyKey.MAX_CHUNKS: 10} - model_type_instance.get_model_schema.return_value = model_schema - - return model_instance - - def test_embedding_dimension_consistency(self, mock_model_instance): - """Test that all embeddings have consistent dimensions. - - Verifies: - - All embeddings have the same dimension - - Dimension matches model specification (1536 for ada-002) - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = [f"Text {i}" for i in range(5)] - - # Create embeddings with consistent dimension - embeddings = [] - for _ in range(5): - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - embeddings.append(normalized) - - usage = EmbeddingUsage( - tokens=50, - total_tokens=50, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.000005"), - currency="USD", - latency=0.7, - ) - - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=embeddings, - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result - - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - assert len(result) == 5 - - # All embeddings should have same dimension - dimensions = [len(emb) for emb in result] - assert all(dim == 1536 for dim in dimensions) - - # All embeddings should be lists of floats - for emb in result: - assert isinstance(emb, list) - assert all(isinstance(x, float) for x in emb) - - def test_embedding_normalization(self, mock_model_instance): - """Test that embeddings are properly normalized (L2 norm ≈ 1.0). - - Verifies: - - All embeddings are L2 normalized - - Normalization is consistent across batches - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = ["Text 1", "Text 2", "Text 3"] - - # Create unnormalized vectors (will be normalized by the service) - embeddings = [] - for _ in range(3): - vector = np.random.randn(1536) * 10 # Unnormalized - normalized = (vector / np.linalg.norm(vector)).tolist() - embeddings.append(normalized) - - usage = EmbeddingUsage( - tokens=30, - total_tokens=30, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.000003"), - currency="USD", - latency=0.5, - ) - - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=embeddings, - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result - - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - for emb in result: - norm = np.linalg.norm(emb) - # L2 norm should be approximately 1.0 - assert abs(norm - 1.0) < 0.01, f"Embedding not normalized: norm={norm}" - - def test_different_model_dimensions(self): - """Test handling of different embedding dimensions for different models. - - Verifies: - - Different models can have different dimensions - - Dimensions are correctly preserved - """ - # Arrange - OpenAI ada-002 (1536 dimensions) - model_instance_ada = Mock() - model_instance_ada.model_name = "text-embedding-ada-002" - model_instance_ada.provider = "openai" - - # Mock model type instance for ada - model_type_instance_ada = Mock() - model_instance_ada.model_type_instance = model_type_instance_ada - model_schema_ada = Mock() - model_schema_ada.model_properties = {ModelPropertyKey.MAX_CHUNKS: 10} - model_type_instance_ada.get_model_schema.return_value = model_schema_ada - - cache_ada = CacheEmbedding(model_instance_ada) - - vector_ada = np.random.randn(1536) - normalized_ada = (vector_ada / np.linalg.norm(vector_ada)).tolist() - - usage_ada = EmbeddingUsage( - tokens=5, - total_tokens=5, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.0000005"), - currency="USD", - latency=0.3, - ) - - result_ada = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=[normalized_ada], - usage=usage_ada, - ) - - # Arrange - Cohere embed-english-v3.0 (1024 dimensions) - model_instance_cohere = Mock() - model_instance_cohere.model_name = "embed-english-v3.0" - model_instance_cohere.provider = "cohere" - - # Mock model type instance for cohere - model_type_instance_cohere = Mock() - model_instance_cohere.model_type_instance = model_type_instance_cohere - model_schema_cohere = Mock() - model_schema_cohere.model_properties = {ModelPropertyKey.MAX_CHUNKS: 10} - model_type_instance_cohere.get_model_schema.return_value = model_schema_cohere - - cache_cohere = CacheEmbedding(model_instance_cohere) - - vector_cohere = np.random.randn(1024) - normalized_cohere = (vector_cohere / np.linalg.norm(vector_cohere)).tolist() - - usage_cohere = EmbeddingUsage( - tokens=5, - total_tokens=5, - unit_price=Decimal("0.0002"), - price_unit=Decimal(1000), - total_price=Decimal("0.000001"), - currency="USD", - latency=0.4, - ) - - result_cohere = EmbeddingResult( - model="embed-english-v3.0", - embeddings=[normalized_cohere], - usage=usage_cohere, - ) - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - - model_instance_ada.invoke_text_embedding.return_value = result_ada - model_instance_cohere.invoke_text_embedding.return_value = result_cohere - - # Act - embedding_ada = cache_ada.embed_documents(["Test"]) - embedding_cohere = cache_cohere.embed_documents(["Test"]) - - # Assert - assert len(embedding_ada[0]) == 1536 # OpenAI dimension - assert len(embedding_cohere[0]) == 1024 # Cohere dimension - - -class TestEmbeddingEdgeCases: - """Test suite for edge cases and special scenarios. - - This class tests unusual inputs and boundary conditions including: - - Empty inputs (empty list, empty strings) - - Very long texts (exceeding typical limits) - - Special characters and Unicode - - Whitespace-only texts - - Duplicate texts in same batch - - Mixed valid and invalid inputs - """ - - @pytest.fixture - def mock_model_instance(self): - """Create a mock ModelInstance for testing. - - Returns: - Mock: Configured ModelInstance with standard settings - - Model: text-embedding-ada-002 - - Provider: openai - - MAX_CHUNKS: 10 - """ - model_instance = Mock() - model_instance.model_name = "text-embedding-ada-002" - model_instance.provider = "openai" - - model_type_instance = Mock() - model_instance.model_type_instance = model_type_instance - - model_schema = Mock() - model_schema.model_properties = {ModelPropertyKey.MAX_CHUNKS: 10} - model_type_instance.get_model_schema.return_value = model_schema - - return model_instance - - def test_embed_empty_list(self, mock_model_instance): - """Test embedding an empty list of documents. - - Verifies: - - Empty list returns empty result - - No model invocation occurs - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = [] - - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - assert result == [] - mock_model_instance.invoke_text_embedding.assert_not_called() - - def test_embed_empty_string(self, mock_model_instance): - """Test embedding an empty string. - - Verifies: - - Empty string is handled correctly - - Model is invoked with empty string - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = [""] - - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - - usage = EmbeddingUsage( - tokens=0, - total_tokens=0, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal(0), - currency="USD", latency=0.1, + ), + ) + + +def _persist_embedding( + session: Session, + *, + text: str, + vector: list[float], + model_name: str = "text-embedding-ada-002", + provider_name: str = "openai", +) -> Embedding: + row = Embedding( + model_name=model_name, + hash=helper.generate_text_hash(text), + provider_name=provider_name, + embedding=pickle.dumps(vector, protocol=pickle.HIGHEST_PROTOCOL), + ) + session.add(row) + session.commit() + return row + + +class TestDocumentCache: + def test_cache_miss_normalizes_and_persists(self, sqlite_embedding_session: Session, model_instance: Mock) -> None: + model_instance.invoke_text_embedding.return_value = _result(_vector()) + + result = CacheEmbedding(model_instance).embed_documents(["new text"]) + + assert np.linalg.norm(result[0]) == pytest.approx(1.0) + model_instance.invoke_text_embedding.assert_called_once_with( + texts=["new text"], input_type=EmbeddingInputType.DOCUMENT ) + persisted = sqlite_embedding_session.scalar(select(Embedding)) + assert persisted is not None + assert persisted.get_embedding() == result[0] - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=[normalized], - usage=usage, - ) + def test_cache_hit_uses_persisted_vector(self, sqlite_embedding_session: Session, model_instance: Mock) -> None: + cached_vector = (_vector() / np.linalg.norm(_vector())).tolist() + row = _persist_embedding(sqlite_embedding_session, text="cached", vector=cached_vector) - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result + result = CacheEmbedding(model_instance).embed_documents(["cached"]) - # Act - result = cache_embedding.embed_documents(texts) + assert result == [cached_vector] + model_instance.invoke_text_embedding.assert_not_called() + assert sqlite_embedding_session.scalars(select(Embedding)).all() == [row] - # Assert - assert len(result) == 1 - assert len(result[0]) == 1536 + def test_partial_hit_preserves_input_order(self, sqlite_embedding_session: Session, model_instance: Mock) -> None: + cached_vector = (_vector(offset=10) / np.linalg.norm(_vector(offset=10))).tolist() + _persist_embedding(sqlite_embedding_session, text="cached", vector=cached_vector) + model_instance.invoke_text_embedding.return_value = _result(_vector(), _vector(offset=20)) - def test_embed_very_long_text(self, mock_model_instance): - """Test embedding very long text. + result = CacheEmbedding(model_instance).embed_documents(["cached", "new one", "new two"]) - Verifies: - - Long texts are handled correctly - - No truncation errors occur - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - # Create a very long text (10000 characters) - long_text = "Python " * 2000 - texts = [long_text] + assert result[0] == cached_vector + assert all(np.linalg.norm(vector) == pytest.approx(1.0) for vector in result) + assert model_instance.invoke_text_embedding.call_args.kwargs["texts"] == ["new one", "new two"] + assert len(sqlite_embedding_session.scalars(select(Embedding)).all()) == 3 - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - - usage = EmbeddingUsage( - tokens=2000, - total_tokens=2000, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.0002"), - currency="USD", - latency=1.5, - ) - - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=[normalized], - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result - - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - assert len(result) == 1 - assert len(result[0]) == 1536 - - def test_embed_special_characters(self, mock_model_instance): - """Test embedding text with special characters. - - Verifies: - - Special characters are handled correctly - - Unicode characters work properly - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = [ - "Hello 世界! 🌍", - "Special chars: @#$%^&*()", - "Newlines\nand\ttabs", + def test_large_batch_respects_model_chunk_limit( + self, sqlite_embedding_session: Session, model_instance: Mock + ) -> None: + texts = [f"text-{index}" for index in range(25)] + model_instance.invoke_text_embedding.side_effect = [ + _result(*(_vector(offset=index) for index in range(10))), + _result(*(_vector(offset=index) for index in range(10, 20))), + _result(*(_vector(offset=index) for index in range(20, 25))), ] - embeddings = [] - for _ in range(3): - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - embeddings.append(normalized) + result = CacheEmbedding(model_instance).embed_documents(texts) - usage = EmbeddingUsage( - tokens=30, - total_tokens=30, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.000003"), - currency="USD", - latency=0.5, - ) - - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=embeddings, - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result - - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - assert len(result) == 3 - assert all(len(emb) == 1536 for emb in result) - - def test_embed_whitespace_only_text(self, mock_model_instance): - """Test embedding text containing only whitespace. - - Verifies: - - Whitespace-only texts are handled correctly - - Model is invoked with whitespace text - - Valid embedding is returned - - Context: - -------- - Whitespace-only texts can occur in real-world scenarios when - processing documents with formatting issues or empty sections. - The embedding model should handle these gracefully. - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = [" ", "\t\t", "\n\n\n"] - - # Create embeddings for whitespace texts - embeddings = [] - for _ in range(3): - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - embeddings.append(normalized) - - usage = EmbeddingUsage( - tokens=3, - total_tokens=3, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.0000003"), - currency="USD", - latency=0.2, - ) - - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=embeddings, - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result - - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - assert len(result) == 3 - assert all(isinstance(emb, list) for emb in result) - assert all(len(emb) == 1536 for emb in result) - - def test_embed_duplicate_texts_in_batch(self, mock_model_instance): - """Test embedding when same text appears multiple times in batch. - - Verifies: - - Duplicate texts are handled correctly - - Each duplicate gets its own embedding - - All duplicates are processed - - Context: - -------- - In batch processing, the same text might appear multiple times. - The current implementation processes all texts individually, - even if they're duplicates. This ensures each position in the - input list gets a corresponding embedding in the output. - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - # Same text repeated 3 times - texts = ["Duplicate text", "Duplicate text", "Duplicate text"] - - # Create embeddings for all three (even though they're duplicates) - embeddings = [] - for _ in range(3): - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - embeddings.append(normalized) - - usage = EmbeddingUsage( - tokens=30, - total_tokens=30, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.000003"), - currency="USD", - latency=0.3, - ) - - # Model returns embeddings for all texts - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=embeddings, - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result - - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - # All three should have embeddings - assert len(result) == 3 - # Model should be called once - mock_model_instance.invoke_text_embedding.assert_called_once() - # All three texts are sent to model (no deduplication) - call_args = mock_model_instance.invoke_text_embedding.call_args - assert len(call_args.kwargs["texts"]) == 3 - - def test_embed_mixed_languages(self, mock_model_instance): - """Test embedding texts in different languages. - - Verifies: - - Multi-language texts are handled correctly - - Unicode characters from various scripts work - - Embeddings are generated for all languages - - Context: - -------- - Modern embedding models support multiple languages. - This test ensures the service handles various scripts: - - Latin (English) - - CJK (Chinese, Japanese, Korean) - - Cyrillic (Russian) - - Arabic - - Emoji and symbols - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = [ - "Hello World", # English - "你好世界", # Chinese - "こんにちは世界", # Japanese - "Привет мир", # Russian - "مرحبا بالعالم", # Arabic - "🌍🌎🌏", # Emoji + assert len(result) == 25 + assert [len(call.kwargs["texts"]) for call in model_instance.invoke_text_embedding.call_args_list] == [ + 10, + 10, + 5, ] + assert len(sqlite_embedding_session.scalars(select(Embedding)).all()) == 25 - # Create embeddings for each language - embeddings = [] - for _ in range(6): - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - embeddings.append(normalized) + @pytest.mark.parametrize( + "provider_error", + [InvokeConnectionError("offline"), InvokeRateLimitError("limited"), InvokeAuthorizationError("denied")], + ) + def test_provider_error_rolls_back_real_transaction( + self, + sqlite_embedding_session: Session, + model_instance: Mock, + provider_error: Exception, + ) -> None: + model_instance.invoke_text_embedding.side_effect = provider_error - usage = EmbeddingUsage( - tokens=60, - total_tokens=60, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.000006"), - currency="USD", - latency=0.8, + with pytest.raises(type(provider_error)): + CacheEmbedding(model_instance).embed_documents(["failure"]) + + assert not sqlite_embedding_session.in_transaction() + assert sqlite_embedding_session.scalar(select(Embedding)) is None + + def test_integrity_race_rolls_back_without_losing_result( + self, sqlite_embedding_session: Session, model_instance: Mock + ) -> None: + model_instance.invoke_text_embedding.return_value = _result(_vector()) + + def raise_integrity_error(_session: Session) -> None: + raise IntegrityError("duplicate embedding", {}, RuntimeError("unique constraint")) + + event.listen(sqlite_embedding_session, "before_commit", raise_integrity_error) + try: + result = CacheEmbedding(model_instance).embed_documents(["racing insert"]) + finally: + event.remove(sqlite_embedding_session, "before_commit", raise_integrity_error) + + assert np.linalg.norm(result[0]) == pytest.approx(1.0) + assert not sqlite_embedding_session.in_transaction() + assert sqlite_embedding_session.scalar(select(Embedding)) is None + + def test_nan_vector_is_skipped_and_logged( + self, + sqlite_embedding_session: Session, + model_instance: Mock, + caplog: pytest.LogCaptureFixture, + ) -> None: + model_instance.invoke_text_embedding.return_value = _result(np.array([np.nan, 1.0])) + + with caplog.at_level("WARNING", logger="core.rag.embedding.cached_embedding"): + result = CacheEmbedding(model_instance).embed_documents(["invalid"]) + + assert result == [None] + assert "Normalized embedding is nan" in caplog.text + assert sqlite_embedding_session.scalar(select(Embedding)) is None + + def test_duplicate_text_is_cached_once(self, sqlite_embedding_session: Session, model_instance: Mock) -> None: + model_instance.invoke_text_embedding.return_value = _result(_vector(), _vector()) + + result = CacheEmbedding(model_instance).embed_documents(["same", "same"]) + + assert result[0] == result[1] + assert len(sqlite_embedding_session.scalars(select(Embedding)).all()) == 1 + + def test_provider_and_model_are_part_of_cache_identity(self, sqlite_embedding_session: Session) -> None: + first = Mock(model_name="shared-model", provider="provider-a", credentials={}) + first.model_type_instance.get_model_schema.return_value = Mock( + model_properties={ModelPropertyKey.MAX_CHUNKS: 10} ) - - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=embeddings, - usage=usage, + first.invoke_text_embedding.return_value = _result(_vector()) + second = Mock(model_name="shared-model", provider="provider-b", credentials={}) + second.model_type_instance.get_model_schema.return_value = Mock( + model_properties={ModelPropertyKey.MAX_CHUNKS: 10} ) + second.invoke_text_embedding.return_value = _result(_vector(offset=10)) - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result + first_result = CacheEmbedding(first).embed_documents(["same text"]) + second_result = CacheEmbedding(second).embed_documents(["same text"]) - # Act - result = cache_embedding.embed_documents(texts) + assert first_result != second_result + assert len(sqlite_embedding_session.scalars(select(Embedding)).all()) == 2 - # Assert - assert len(result) == 6 - assert all(isinstance(emb, list) for emb in result) - assert all(len(emb) == 1536 for emb in result) - # Verify all embeddings are normalized - for emb in result: - norm = np.linalg.norm(emb) - assert abs(norm - 1.0) < 0.01 + def test_empty_input_does_not_touch_provider(self, model_instance: Mock) -> None: + assert CacheEmbedding(model_instance).embed_documents([]) == [] + model_instance.invoke_text_embedding.assert_not_called() - def test_embed_query_uses_bound_model_instance(self, mock_model_instance): - """Test query embedding using the provided model instance. - Verifies: - - Embedding generation works with the injected model instance - - Query input type is preserved - - No extra binding step is required at call time - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - query = "What is machine learning?" +class TestMultimodalDocumentCache: + def test_multimodal_miss_persists_and_hit_reuses( + self, sqlite_embedding_session: Session, model_instance: Mock + ) -> None: + document = {"file_id": "file-1"} + model_instance.invoke_multimodal_embedding.return_value = _result(_vector()) + cache = CacheEmbedding(model_instance) - # Create embedding - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() + generated = cache.embed_multimodal_documents([document]) + cached = cache.embed_multimodal_documents([document]) - usage = EmbeddingUsage( - tokens=5, - total_tokens=5, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.0000005"), - currency="USD", - latency=0.3, + assert cached == generated + model_instance.invoke_multimodal_embedding.assert_called_once_with( + multimodel_documents=[document], input_type=EmbeddingInputType.DOCUMENT ) + assert sqlite_embedding_session.scalar(select(Embedding).where(Embedding.hash == "file-1")) is not None - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=[normalized], - usage=usage, + +class TestQueryCache: + @patch("core.rag.embedding.cached_embedding.redis_client") + def test_query_cache_miss_normalizes_and_stores(self, redis: Mock, model_instance: Mock) -> None: + redis.get.return_value = None + model_instance.invoke_text_embedding.return_value = _result(_vector()) + + result = CacheEmbedding(model_instance).embed_query("query") + + assert np.linalg.norm(result) == pytest.approx(1.0) + model_instance.invoke_text_embedding.assert_called_once_with( + texts=["query"], input_type=EmbeddingInputType.QUERY ) + redis.setex.assert_called_once() - with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: - mock_redis.get.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result + @patch("core.rag.embedding.cached_embedding.redis_client") + def test_query_cache_hit_refreshes_ttl(self, redis: Mock, model_instance: Mock) -> None: + vector = np.array([0.25, 0.75], dtype=float) + redis.get.return_value = base64.b64encode(vector.tobytes()) - # Act - result = cache_embedding.embed_query(query) + result = CacheEmbedding(model_instance).embed_query("query") - # Assert - assert isinstance(result, list) - assert len(result) == 1536 + assert result == vector.tolist() + redis.expire.assert_called_once_with(f"openai_text-embedding-ada-002_{helper.generate_text_hash('query')}", 600) + model_instance.invoke_text_embedding.assert_not_called() - mock_model_instance.invoke_text_embedding.assert_called_once_with( - texts=[query], - input_type=EmbeddingInputType.QUERY, - ) + @patch("core.rag.embedding.cached_embedding.redis_client") + def test_query_nan_is_rejected(self, redis: Mock, model_instance: Mock) -> None: + redis.get.return_value = None + model_instance.invoke_text_embedding.return_value = _result(np.array([np.nan, 1.0])) - def test_embed_documents_uses_bound_model_instance(self, mock_model_instance): - """Test document embedding using the provided model instance. + with pytest.raises(ValueError, match="Normalized embedding is nan"): + CacheEmbedding(model_instance).embed_query("query") - Verifies: - - Batch processing uses the injected model instance - - Document input type is preserved - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - texts = ["Document 1", "Document 2"] + redis.setex.assert_not_called() - # Create embeddings - embeddings = [] - for _ in range(2): - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - embeddings.append(normalized) + @patch("core.rag.embedding.cached_embedding.redis_client") + def test_redis_write_error_is_propagated(self, redis: Mock, model_instance: Mock) -> None: + redis.get.return_value = None + redis.setex.side_effect = ConnectionError("redis unavailable") + model_instance.invoke_text_embedding.return_value = _result(_vector()) - usage = EmbeddingUsage( - tokens=20, - total_tokens=20, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.000002"), - currency="USD", - latency=0.5, - ) + with pytest.raises(ConnectionError, match="redis unavailable"): + CacheEmbedding(model_instance).embed_query("query") - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=embeddings, - usage=usage, - ) + @patch("core.rag.embedding.cached_embedding.redis_client") + def test_multimodal_query_uses_file_cache_key(self, redis: Mock, model_instance: Mock) -> None: + redis.get.return_value = None + model_instance.invoke_multimodal_embedding.return_value = _result(_vector()) - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result + result = CacheEmbedding(model_instance).embed_multimodal_query({"file_id": "file-1"}) - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - assert len(result) == 2 - - mock_model_instance.invoke_text_embedding.assert_called_once() - call_args = mock_model_instance.invoke_text_embedding.call_args - assert call_args.kwargs["input_type"] == EmbeddingInputType.DOCUMENT - - -class TestEmbeddingCachePerformance: - """Test suite for cache performance and optimization scenarios. - - This class tests cache-related performance optimizations: - - Cache hit rate improvements - - Batch processing efficiency - - Memory usage optimization - - Cache key generation - - TTL (Time To Live) management - """ - - @pytest.fixture - def mock_model_instance(self): - """Create a mock ModelInstance for testing. - - Returns: - Mock: Configured ModelInstance for performance testing - - Model: text-embedding-ada-002 - - Provider: openai - - MAX_CHUNKS: 10 - """ - model_instance = Mock() - model_instance.model_name = "text-embedding-ada-002" - model_instance.provider = "openai" - - model_type_instance = Mock() - model_instance.model_type_instance = model_type_instance - - model_schema = Mock() - model_schema.model_properties = {ModelPropertyKey.MAX_CHUNKS: 10} - model_type_instance.get_model_schema.return_value = model_schema - - return model_instance - - def test_cache_hit_reduces_api_calls(self, mock_model_instance): - """Test that cache hits prevent unnecessary API calls. - - Verifies: - - First call triggers API request - - Second call uses cache (no API call) - - Cache significantly reduces API usage - - Context: - -------- - Caching is critical for: - 1. Reducing API costs - 2. Improving response time - 3. Reducing rate limit pressure - 4. Better user experience - - This test demonstrates the cache working as expected. - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - text = "Frequently used text" - - # Create cached embedding - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - - mock_cached_embedding = Mock(spec=Embedding) - mock_cached_embedding.get_embedding.return_value = normalized - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - # First call: cache miss - mock_session.scalar.return_value = None - - usage = EmbeddingUsage( - tokens=5, - total_tokens=5, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.0000005"), - currency="USD", - latency=0.3, - ) - - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=[normalized], - usage=usage, - ) - - mock_model_instance.invoke_text_embedding.return_value = embedding_result - - # Act - First call (cache miss) - result1 = cache_embedding.embed_documents([text]) - - # Assert - Model was called - assert mock_model_instance.invoke_text_embedding.call_count == 1 - assert len(result1) == 1 - - # Arrange - Second call: cache hit - mock_session.scalar.return_value = mock_cached_embedding - - # Act - Second call (cache hit) - result2 = cache_embedding.embed_documents([text]) - - # Assert - Model was NOT called again (still 1 call total) - assert mock_model_instance.invoke_text_embedding.call_count == 1 - assert len(result2) == 1 - assert result2[0] == normalized # Same embedding from cache - - def test_batch_processing_efficiency(self, mock_model_instance): - """Test that batch processing is more efficient than individual calls. - - Verifies: - - Multiple texts are processed in single API call - - Batch size respects MAX_CHUNKS limit - - Batching reduces total API calls - - Context: - -------- - Batch processing is essential for: - 1. Reducing API overhead - 2. Better throughput - 3. Lower latency per text - 4. Cost optimization - - Example: 100 texts in batches of 10 = 10 API calls - vs 100 individual calls = 100 API calls - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - # 15 texts should be processed in 2 batches (10 + 5) - texts = [f"Text {i}" for i in range(15)] - - # Create embeddings for each batch - def create_batch_result(batch_size): - """Helper function to create batch embedding results.""" - embeddings = [] - for _ in range(batch_size): - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - embeddings.append(normalized) - - usage = EmbeddingUsage( - tokens=batch_size * 10, - total_tokens=batch_size * 10, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal(str(batch_size * 0.000001)), - currency="USD", - latency=0.5, - ) - - return EmbeddingResult( - model="text-embedding-ada-002", - embeddings=embeddings, - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.db.session") as mock_session: - mock_session.scalar.return_value = None - - # Mock model to return appropriate batch results - batch_results = [ - create_batch_result(10), # First batch - create_batch_result(5), # Second batch - ] - mock_model_instance.invoke_text_embedding.side_effect = batch_results - - # Act - result = cache_embedding.embed_documents(texts) - - # Assert - assert len(result) == 15 - # Only 2 API calls for 15 texts (batched) - assert mock_model_instance.invoke_text_embedding.call_count == 2 - - # Verify batch sizes - calls = mock_model_instance.invoke_text_embedding.call_args_list - assert len(calls[0].kwargs["texts"]) == 10 # First batch - assert len(calls[1].kwargs["texts"]) == 5 # Second batch - - def test_redis_cache_expiration(self, mock_model_instance): - """Test Redis cache TTL (Time To Live) management. - - Verifies: - - Cache entries have appropriate TTL (600 seconds) - - TTL is extended on cache hits - - Expired entries are regenerated - - Context: - -------- - Redis cache TTL ensures: - 1. Memory doesn't grow unbounded - 2. Stale embeddings are refreshed - 3. Frequently used queries stay cached longer - 4. Infrequently used queries expire naturally - """ - # Arrange - cache_embedding = CacheEmbedding(mock_model_instance) - query = "Test query" - - vector = np.random.randn(1536) - normalized = (vector / np.linalg.norm(vector)).tolist() - - usage = EmbeddingUsage( - tokens=5, - total_tokens=5, - unit_price=Decimal("0.0001"), - price_unit=Decimal(1000), - total_price=Decimal("0.0000005"), - currency="USD", - latency=0.3, - ) - - embedding_result = EmbeddingResult( - model="text-embedding-ada-002", - embeddings=[normalized], - usage=usage, - ) - - with patch("core.rag.embedding.cached_embedding.redis_client") as mock_redis: - # Test cache miss - sets TTL - mock_redis.get.return_value = None - mock_model_instance.invoke_text_embedding.return_value = embedding_result - - # Act - cache_embedding.embed_query(query) - - # Assert - TTL was set to 600 seconds - mock_redis.setex.assert_called_once() - call_args = mock_redis.setex.call_args - assert call_args[0][1] == 600 # TTL in seconds - - # Test cache hit - extends TTL - mock_redis.reset_mock() - vector_bytes = np.array(normalized).tobytes() - encoded_vector = base64.b64encode(vector_bytes).decode("utf-8") - mock_redis.get.return_value = encoded_vector - - # Act - cache_embedding.embed_query(query) - - # Assert - TTL was extended - mock_redis.expire.assert_called_once() - assert mock_redis.expire.call_args[0][1] == 600 + assert np.linalg.norm(result) == pytest.approx(1.0) + assert redis.setex.call_args.args[:2] == ("openai_text-embedding-ada-002_file-1", 600)