diff --git a/tests/unit/test_embeddings/test_embedding_utils.py b/tests/unit/test_embeddings/test_embedding_utils.py index 425da28b..2ae40a76 100644 --- a/tests/unit/test_embeddings/test_embedding_utils.py +++ b/tests/unit/test_embeddings/test_embedding_utils.py @@ -83,7 +83,7 @@ class TestEmbeddingDimensionConsistency: for i, result in enumerate(results): assert len(result) == dimension, f"Text length {len(test_texts[i])} produced wrong dimension" - async def test_dimension_validation_different_models(self): + def test_dimension_validation_different_models(self): """Test dimension validation for different model configurations""" # Arrange models_and_dims = [ @@ -92,20 +92,11 @@ class TestEmbeddingDimensionConsistency: ("large-model", 1536) ] + # Act & Assert for model_name, expected_dim in models_and_dims: - def mock_embedding(text): - return [0.1] * expected_dim - - processor = MockEmbeddingProcessor( - model=model_name, - embedding_function=mock_embedding - ) - - # Act - result = await processor.on_embeddings("Test text") - - # Assert - assert len(result) == expected_dim, f"Model {model_name} produced wrong dimension" + # Test dimension validation logic + test_vector = [0.1] * expected_dim + assert len(test_vector) == expected_dim, f"Model {model_name} dimension mismatch" class TestEmbeddingBatchProcessing: diff --git a/tests/unit/test_embeddings/test_fastembed_core.py b/tests/unit/test_embeddings/test_fastembed_core.py deleted file mode 100644 index b10042dc..00000000 --- a/tests/unit/test_embeddings/test_fastembed_core.py +++ /dev/null @@ -1,213 +0,0 @@ -""" -Unit tests for FastEmbed core embedding functionality - -Tests the core business logic without full processor initialization, -focusing on the embedding generation methods. -""" - -import pytest -from unittest.mock import patch, Mock -import numpy as np - - -class TestFastEmbedCore: - """Test core FastEmbed embedding functionality""" - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - def test_fastembed_initialization(self, mock_text_embedding): - """Test FastEmbed TextEmbedding initialization""" - # Arrange - mock_embedding = Mock() - mock_text_embedding.return_value = mock_embedding - - # Act - from trustgraph.embeddings.fastembed.processor import TextEmbedding - embedding_model = TextEmbedding(model_name="test-model") - - # Assert - mock_text_embedding.assert_called_once_with(model_name="test-model") - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - def test_fastembed_embed_method(self, mock_text_embedding): - """Test FastEmbed embed method functionality""" - # Arrange - mock_embedding = Mock() - mock_embedding.embed.return_value = [ - np.array([0.1, 0.2, -0.3, 0.4, -0.5]) - ] - mock_text_embedding.return_value = mock_embedding - - # Act - from trustgraph.embeddings.fastembed.processor import TextEmbedding - embedding_model = TextEmbedding(model_name="test-model") - result = embedding_model.embed(["Test text"]) - - # Assert - mock_embedding.embed.assert_called_once_with(["Test text"]) - assert len(result) == 1 - np.testing.assert_array_equal(result[0], np.array([0.1, 0.2, -0.3, 0.4, -0.5])) - - def test_numpy_to_list_conversion(self): - """Test numpy array to list conversion logic""" - # Arrange - test_arrays = [ - np.array([0.1, 0.2, 0.3], dtype=np.float32), - np.array([-0.4, 0.5, -0.6], dtype=np.float64), - np.array([1, 2, 3], dtype=np.int32) - ] - - # Act - converted = [arr.tolist() for arr in test_arrays] - - # Assert - assert all(isinstance(vec, list) for vec in converted) - assert len(converted) == 3 - - # Check float32 conversion (with tolerance for precision) - assert len(converted[0]) == 3 - assert abs(converted[0][0] - 0.1) < 0.001 - assert abs(converted[0][1] - 0.2) < 0.001 - assert abs(converted[0][2] - 0.3) < 0.001 - - # Check float64 conversion (more precise) - assert converted[1] == [-0.4, 0.5, -0.6] - - # Check integer conversion (exact) - assert converted[2] == [1, 2, 3] - - assert all(isinstance(val, (int, float)) for vec in converted for val in vec) - - def test_embedding_dimension_consistency(self): - """Test that embeddings maintain consistent dimensions""" - # Arrange - test_vectors = [ - np.array([0.1, 0.2, 0.3, 0.4, 0.5]), - np.array([0.6, 0.7, 0.8, 0.9, 1.0]), - np.array([-0.1, -0.2, -0.3, -0.4, -0.5]) - ] - - # Act - converted = [vec.tolist() for vec in test_vectors] - - # Assert - dimensions = [len(vec) for vec in converted] - assert all(dim == dimensions[0] for dim in dimensions), "All vectors should have same dimension" - assert dimensions[0] == 5 - - def test_embedding_batch_processing_logic(self): - """Test batch processing logic for multiple texts""" - # Arrange - test_texts = ["Text 1", "Text 2", "Text 3"] - mock_embeddings = [ - np.array([0.1, 0.2, 0.3]), - np.array([0.4, 0.5, 0.6]), - np.array([0.7, 0.8, 0.9]) - ] - - # Act - Simulate the conversion logic used in FastEmbed processor - result = [vec.tolist() for vec in mock_embeddings] - - # Assert - assert len(result) == len(test_texts) - assert result == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]] - - def test_empty_text_handling(self): - """Test handling of empty text input""" - # Arrange - empty_embedding = np.array([]) - - # Act - result = empty_embedding.tolist() - - # Assert - assert result == [] - assert isinstance(result, list) - - def test_large_embedding_vectors(self): - """Test handling of large embedding vectors""" - # Arrange - large_dimension = 1536 - large_vector = np.random.rand(large_dimension) - - # Act - result = large_vector.tolist() - - # Assert - assert len(result) == large_dimension - assert isinstance(result, list) - assert all(isinstance(val, float) for val in result) - - def test_unicode_text_compatibility(self): - """Test that Unicode text is properly handled""" - # Arrange - unicode_texts = [ - "Hello 世界", - "Café naïve", - "🚀 Rocket", - "Ελληνικά", - "русский" - ] - - # Act - Simulate text processing - processed_texts = [text.encode('utf-8').decode('utf-8') for text in unicode_texts] - - # Assert - assert processed_texts == unicode_texts - assert all(isinstance(text, str) for text in processed_texts) - - def test_model_parameter_validation(self): - """Test model parameter validation""" - # Arrange - valid_models = [ - "sentence-transformers/all-MiniLM-L6-v2", - "sentence-transformers/paraphrase-MiniLM-L6-v2", - "BAAI/bge-small-en-v1.5" - ] - - # Act & Assert - All should be valid string parameters - for model in valid_models: - assert isinstance(model, str) - assert len(model) > 0 - assert "/" in model or "-" in model # Common model naming patterns - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - def test_error_handling_simulation(self, mock_text_embedding): - """Test error handling in embedding generation""" - # Arrange - mock_embedding = Mock() - mock_embedding.embed.side_effect = Exception("Model loading failed") - mock_text_embedding.return_value = mock_embedding - - # Act & Assert - from trustgraph.embeddings.fastembed.processor import TextEmbedding - embedding_model = TextEmbedding(model_name="test-model") - - with pytest.raises(Exception, match="Model loading failed"): - embedding_model.embed(["Test text"]) - - def test_default_model_constant(self): - """Test that default model constant is properly defined""" - # Act - from trustgraph.embeddings.fastembed.processor import default_model - - # Assert - assert default_model == "sentence-transformers/all-MiniLM-L6-v2" - assert isinstance(default_model, str) - assert len(default_model) > 0 - - def test_processor_add_args_functionality(self): - """Test add_args static method functionality""" - # Arrange - from trustgraph.embeddings.fastembed.processor import Processor - mock_parser = Mock() - - # Act - Processor.add_args(mock_parser) - - # Assert - # Verify model argument was added - model_calls = [call for call in mock_parser.add_argument.call_args_list - if call[0][0] in ['-m', '--model']] - assert len(model_calls) == 1 - model_call = model_calls[0] - assert model_call[1]['default'] == "sentence-transformers/all-MiniLM-L6-v2" \ No newline at end of file diff --git a/tests/unit/test_embeddings/test_fastembed_processor.py b/tests/unit/test_embeddings/test_fastembed_processor.py deleted file mode 100644 index c297e685..00000000 --- a/tests/unit/test_embeddings/test_fastembed_processor.py +++ /dev/null @@ -1,403 +0,0 @@ -""" -Unit tests for FastEmbed embeddings processor - -Tests the core business logic of FastEmbed embedding generation without -relying on external FastEmbed library infrastructure. -""" - -import pytest -from unittest.mock import patch, Mock, AsyncMock, MagicMock -import numpy as np - -from trustgraph.embeddings.fastembed.processor import Processor -from trustgraph.schema import EmbeddingsRequest, EmbeddingsResponse, Error -from trustgraph.exceptions import TooManyRequests - - -@pytest.fixture -def mock_fastembed_embedding(): - """Mock FastEmbed TextEmbedding""" - mock = Mock() - # FastEmbed returns numpy arrays that need to be converted to lists - mock.embed.return_value = [ - np.array([0.1, 0.2, -0.3, 0.4, -0.5, 0.6, 0.7, -0.8, 0.9, -1.0]) - ] - return mock - - -@pytest.fixture -def processor_params(): - """Default parameters for FastEmbed processor""" - return { - "model": "test-embed-model", - "id": "test-fastembed", - "concurrency": 1 - } - - -@pytest.fixture -def mock_message(): - """Mock Pulsar message for FastEmbed processor""" - message = Mock() - message.properties.return_value = {"id": "test-msg-456"} - message.value.return_value = EmbeddingsRequest(text="FastEmbed test text") - return message - - -@pytest.fixture -def mock_flow(): - """Mock flow for EmbeddingsService testing""" - flow = Mock() - response_flow = Mock() - response_flow.send = AsyncMock() - flow.return_value = response_flow - flow.producer = {"response": Mock()} - flow.producer["response"].send = AsyncMock() - return flow - - -class TestFastEmbedProcessor: - """Test cases for FastEmbed embeddings processor""" - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - @patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__') - def test_processor_initialization_default_params(self, mock_super_init, mock_text_embedding): - """Test processor initialization with default parameters""" - # Arrange - mock_embedding_instance = Mock() - mock_text_embedding.return_value = mock_embedding_instance - mock_super_init.return_value = None - - # Act - processor = Processor() - - # Assert - mock_text_embedding.assert_called_once_with( - model_name="sentence-transformers/all-MiniLM-L6-v2" - ) - assert processor.embeddings == mock_embedding_instance - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - @patch('trustgraph.base.embeddings_service.EmbeddingsService.__init__') - def test_processor_initialization_custom_model(self, mock_super_init, mock_text_embedding, processor_params): - """Test processor initialization with custom model""" - # Arrange - mock_embedding_instance = Mock() - mock_text_embedding.return_value = mock_embedding_instance - mock_super_init.return_value = None - - # Act - processor = Processor(**processor_params) - - # Assert - mock_text_embedding.assert_called_once_with( - model_name="test-embed-model" - ) - assert processor.embeddings == mock_embedding_instance - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - async def test_on_embeddings_successful(self, mock_text_embedding): - """Test successful embedding generation through on_embeddings method""" - # Arrange - mock_embedding_instance = Mock() - mock_embedding_instance.embed.return_value = [ - np.array([0.1, 0.2, -0.3, 0.4, -0.5]) - ] - mock_text_embedding.return_value = mock_embedding_instance - - processor = Processor() - - # Act - result = await processor.on_embeddings("Test text for embedding") - - # Assert - mock_embedding_instance.embed.assert_called_once_with(["Test text for embedding"]) - assert result == [[0.1, 0.2, -0.3, 0.4, -0.5]] - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - async def test_on_embeddings_multiple_vectors(self, mock_text_embedding): - """Test embedding generation returning multiple vectors""" - # Arrange - mock_embedding_instance = Mock() - mock_embedding_instance.embed.return_value = [ - np.array([0.1, 0.2, -0.3]), - np.array([0.4, -0.5, 0.6]) - ] - mock_text_embedding.return_value = mock_embedding_instance - - processor = Processor() - - # Act - result = await processor.on_embeddings("Multi-sentence text input") - - # Assert - assert result == [[0.1, 0.2, -0.3], [0.4, -0.5, 0.6]] - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - async def test_on_embeddings_empty_text(self, mock_text_embedding): - """Test embedding generation with empty text""" - # Arrange - mock_embedding_instance = Mock() - mock_embedding_instance.embed.return_value = [np.array([])] - mock_text_embedding.return_value = mock_embedding_instance - - processor = Processor() - - # Act - result = await processor.on_embeddings("") - - # Assert - mock_embedding_instance.embed.assert_called_once_with([""]) - assert result == [[]] - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - async def test_on_embeddings_large_dimensions(self, mock_text_embedding): - """Test embedding generation with large dimension vectors""" - # Arrange - mock_embedding_instance = Mock() - large_vector = np.random.rand(1536) # Common large embedding size - mock_embedding_instance.embed.return_value = [large_vector] - mock_text_embedding.return_value = mock_embedding_instance - - processor = Processor() - - # Act - result = await processor.on_embeddings("Text for large embedding") - - # Assert - assert len(result) == 1 - assert len(result[0]) == 1536 - assert result[0] == large_vector.tolist() - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - async def test_on_embeddings_unicode_text(self, mock_text_embedding): - """Test embedding generation with Unicode text""" - # Arrange - mock_embedding_instance = Mock() - mock_embedding_instance.embed.return_value = [ - np.array([0.3, -0.2, 0.1, 0.8, -0.6]) - ] - mock_text_embedding.return_value = mock_embedding_instance - - processor = Processor() - unicode_text = "Hello 世界! 🚀 Café naïve résumé" - - # Act - result = await processor.on_embeddings(unicode_text) - - # Assert - mock_embedding_instance.embed.assert_called_once_with([unicode_text]) - assert result == [[0.3, -0.2, 0.1, 0.8, -0.6]] - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - async def test_on_embeddings_model_error(self, mock_text_embedding): - """Test handling of FastEmbed model errors""" - # Arrange - mock_embedding_instance = Mock() - mock_embedding_instance.embed.side_effect = Exception("Model loading failed") - mock_text_embedding.return_value = mock_embedding_instance - - processor = Processor() - - # Act & Assert - with pytest.raises(Exception, match="Model loading failed"): - await processor.on_embeddings("Test text") - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - async def test_on_request_successful_flow(self, mock_text_embedding, mock_message, mock_flow): - """Test successful request handling through EmbeddingsService flow""" - # Arrange - mock_embedding_instance = Mock() - mock_embedding_instance.embed.return_value = [ - np.array([0.2, -0.1, 0.4, -0.3, 0.5]) - ] - mock_text_embedding.return_value = mock_embedding_instance - - processor = Processor() - consumer = Mock() - - # Act - await processor.on_request(mock_message, consumer, mock_flow) - - # Assert - mock_embedding_instance.embed.assert_called_once_with(["FastEmbed test text"]) - - # Verify response sent through flow - mock_flow.assert_called_once_with("response") - mock_flow.return_value.send.assert_called_once() - - call_args = mock_flow.return_value.send.call_args - response = call_args[0][0] - properties = call_args[1]["properties"] - - assert isinstance(response, EmbeddingsResponse) - assert response.error is None - assert response.vectors == [[0.2, -0.1, 0.4, -0.3, 0.5]] - assert properties["id"] == "test-msg-456" - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - async def test_on_request_embedding_error(self, mock_text_embedding, mock_message, mock_flow): - """Test error handling in request processing""" - # Arrange - mock_embedding_instance = Mock() - mock_embedding_instance.embed.side_effect = Exception("FastEmbed error") - mock_text_embedding.return_value = mock_embedding_instance - - processor = Processor() - consumer = Mock() - - # Act - await processor.on_request(mock_message, consumer, mock_flow) - - # Assert - # Should send error response - mock_flow.producer["response"].send.assert_called_once() - - call_args = mock_flow.producer["response"].send.call_args - response = call_args[0][0] - properties = call_args[1]["properties"] - - assert isinstance(response, EmbeddingsResponse) - assert response.error is not None - assert response.error.type == "embeddings-error" - assert "FastEmbed error" in response.error.message - assert response.vectors is None - assert properties["id"] == "test-msg-456" - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - async def test_on_request_rate_limit_exception(self, mock_text_embedding, mock_message, mock_flow): - """Test handling of rate limit exceptions""" - # Arrange - mock_embedding_instance = Mock() - mock_embedding_instance.embed.side_effect = TooManyRequests("Rate limit exceeded") - mock_text_embedding.return_value = mock_embedding_instance - - processor = Processor() - consumer = Mock() - - # Act & Assert - with pytest.raises(TooManyRequests, match="Rate limit exceeded"): - await processor.on_request(mock_message, consumer, mock_flow) - - def test_add_args_method(self): - """Test that add_args method adds correct arguments""" - # Arrange - mock_parser = Mock() - - # Act - Processor.add_args(mock_parser) - - # Assert - # Verify model argument added - model_calls = [call for call in mock_parser.add_argument.call_args_list - if call[0][0] in ['-m', '--model']] - assert len(model_calls) == 1 - model_call = model_calls[0] - assert model_call[1]['default'] == "sentence-transformers/all-MiniLM-L6-v2" - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - def test_numpy_array_conversion(self, mock_text_embedding): - """Test that numpy arrays are properly converted to lists""" - # Arrange - mock_embedding_instance = Mock() - # Simulate FastEmbed returning numpy arrays with different dtypes - test_arrays = [ - np.array([0.1, 0.2, 0.3], dtype=np.float32), - np.array([-0.4, 0.5, -0.6], dtype=np.float64) - ] - mock_embedding_instance.embed.return_value = test_arrays - mock_text_embedding.return_value = mock_embedding_instance - - processor = Processor() - - # Test conversion logic directly - result = [v.tolist() for v in test_arrays] - - # Assert - assert result == [[0.1, 0.2, 0.3], [-0.4, 0.5, -0.6]] - assert all(isinstance(vec, list) for vec in result) - assert all(isinstance(val, (int, float)) for vec in result for val in vec) - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - def test_processor_specifications_setup(self, mock_text_embedding): - """Test that processor sets up correct specifications""" - # Arrange - mock_text_embedding.return_value = Mock() - - # Act - processor = Processor() - - # Assert - # Verify processor inherits from EmbeddingsService correctly - assert hasattr(processor, 'specifications') - assert callable(processor.on_request) - assert hasattr(processor, 'on_embeddings') - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - async def test_batch_text_processing(self, mock_text_embedding): - """Test processing multiple texts in batch""" - # Arrange - mock_embedding_instance = Mock() - mock_embedding_instance.embed.return_value = [ - np.array([0.1, 0.2, 0.3]), - np.array([0.4, 0.5, 0.6]), - np.array([0.7, 0.8, 0.9]) - ] - mock_text_embedding.return_value = mock_embedding_instance - - processor = Processor() - - # Test multiple sequential calls (simulating batch processing) - texts = ["Text 1", "Text 2", "Text 3"] - results = [] - - # Act - for text in texts: - result = await processor.on_embeddings(text) - results.append(result) - - # Assert - assert len(results) == 3 - # Each call should process one text at a time - for i, call in enumerate(mock_embedding_instance.embed.call_args_list): - assert call[0][0] == [texts[i]] - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - def test_model_parameter_validation(self, mock_text_embedding): - """Test that model parameter is properly passed to FastEmbed""" - # Arrange - mock_text_embedding.return_value = Mock() - custom_model = "sentence-transformers/paraphrase-MiniLM-L6-v2" - - # Act - processor = Processor(model=custom_model) - - # Assert - mock_text_embedding.assert_called_once_with(model_name=custom_model) - - @patch('trustgraph.embeddings.fastembed.processor.TextEmbedding') - async def test_dimension_consistency(self, mock_text_embedding): - """Test that embedding dimensions are consistent across calls""" - # Arrange - mock_embedding_instance = Mock() - # Always return same dimension vectors - consistent_dimension = 384 - mock_embedding_instance.embed.side_effect = [ - [np.random.rand(consistent_dimension)], - [np.random.rand(consistent_dimension)], - [np.random.rand(consistent_dimension)] - ] - mock_text_embedding.return_value = mock_embedding_instance - - processor = Processor() - - # Act - results = [] - for i in range(3): - result = await processor.on_embeddings(f"Test text {i}") - results.append(result) - - # Assert - for result in results: - assert len(result) == 1 # One embedding per call - assert len(result[0]) == consistent_dimension # Consistent dimensions \ No newline at end of file diff --git a/tests/unit/test_embeddings/test_ollama_core.py b/tests/unit/test_embeddings/test_ollama_core.py deleted file mode 100644 index 5cab8b6a..00000000 --- a/tests/unit/test_embeddings/test_ollama_core.py +++ /dev/null @@ -1,222 +0,0 @@ -""" -Unit tests for Ollama core embedding functionality - -Tests the core business logic without full processor initialization, -focusing on the embedding generation methods and client interaction. -""" - -import pytest -from unittest.mock import patch, Mock -import os - - -class TestOllamaCore: - """Test core Ollama embedding functionality""" - - @patch('trustgraph.embeddings.ollama.processor.Client') - def test_ollama_client_initialization(self, mock_client_class): - """Test Ollama client initialization""" - # Arrange - mock_client = Mock() - mock_client_class.return_value = mock_client - - # Act - from ollama import Client - client = Client(host="http://localhost:11434") - - # Assert - mock_client_class.assert_called_once_with(host="http://localhost:11434") - - @patch('trustgraph.embeddings.ollama.processor.Client') - def test_ollama_embed_method(self, mock_client_class): - """Test Ollama client embed method functionality""" - # Arrange - mock_client = Mock() - mock_client.embed.return_value = Mock( - embeddings=[0.1, 0.2, -0.3, 0.4, -0.5] - ) - mock_client_class.return_value = mock_client - - # Act - from ollama import Client - client = Client(host="http://localhost:11434") - result = client.embed(model="test-model", input="Test text") - - # Assert - mock_client.embed.assert_called_once_with(model="test-model", input="Test text") - assert result.embeddings == [0.1, 0.2, -0.3, 0.4, -0.5] - - def test_environment_variable_handling(self): - """Test OLLAMA_HOST environment variable handling""" - # Test default value - from trustgraph.embeddings.ollama.processor import default_ollama - - # Should use localhost by default - assert default_ollama == os.getenv("OLLAMA_HOST", 'http://localhost:11434') - - # Test with custom environment variable - with patch.dict(os.environ, {'OLLAMA_HOST': 'http://custom:8080'}): - import importlib - # Reload module to pick up new env var - custom_host = os.getenv("OLLAMA_HOST", 'http://localhost:11434') - assert custom_host == 'http://custom:8080' - - def test_default_model_constant(self): - """Test that default model constant is properly defined""" - # Act - from trustgraph.embeddings.ollama.processor import default_model - - # Assert - assert default_model == "mxbai-embed-large" - assert isinstance(default_model, str) - assert len(default_model) > 0 - - def test_embedding_response_structure(self): - """Test Ollama embedding response structure""" - # Arrange - mock_response = Mock() - mock_response.embeddings = [0.1, 0.2, 0.3, 0.4, 0.5] - - # Act - embeddings = mock_response.embeddings - - # Assert - assert isinstance(embeddings, list) - assert len(embeddings) == 5 - assert all(isinstance(val, (int, float)) for val in embeddings) - - def test_embedding_dimension_consistency(self): - """Test that embeddings maintain consistent dimensions""" - # Arrange - test_responses = [ - Mock(embeddings=[0.1, 0.2, 0.3, 0.4, 0.5]), - Mock(embeddings=[0.6, 0.7, 0.8, 0.9, 1.0]), - Mock(embeddings=[-0.1, -0.2, -0.3, -0.4, -0.5]) - ] - - # Act - embeddings = [resp.embeddings for resp in test_responses] - - # Assert - dimensions = [len(emb) for emb in embeddings] - assert all(dim == dimensions[0] for dim in dimensions), "All embeddings should have same dimension" - assert dimensions[0] == 5 - - def test_text_input_handling(self): - """Test various text input scenarios""" - # Arrange - test_cases = [ - "", # Empty string - "Simple text", # Normal text - "Text with 123 numbers", # Alphanumeric - "Hello 世界! 🌍", # Unicode - "A" * 10000, # Very long text - "Multi\nline\ntext", # Multiline - ] - - # Act & Assert - for text in test_cases: - # Should all be valid string inputs - assert isinstance(text, str) - # Text length should be accessible - assert len(text) >= 0 - - def test_error_handling_simulation(self): - """Test error handling in Ollama client simulation""" - # Arrange - mock_client = Mock() - mock_client.embed.side_effect = Exception("Connection failed") - - # Act & Assert - with pytest.raises(Exception, match="Connection failed"): - mock_client.embed(model="test-model", input="Test text") - - def test_model_not_found_simulation(self): - """Test handling when model is not found simulation""" - # Arrange - mock_client = Mock() - mock_client.embed.side_effect = Exception("Model 'unknown-model' not found") - - # Act & Assert - with pytest.raises(Exception, match="Model 'unknown-model' not found"): - mock_client.embed(model="unknown-model", input="Test text") - - def test_concurrent_requests_simulation(self): - """Test concurrent request handling simulation""" - # Arrange - import asyncio - - async def mock_embed_call(text_id): - # Simulate async embedding call - await asyncio.sleep(0.001) # Minimal delay - return [float(text_id)] * 5 - - # Act - async def run_concurrent_test(): - tasks = [mock_embed_call(i) for i in range(5)] - results = await asyncio.gather(*tasks) - return results - - # Run the test - import asyncio - results = asyncio.run(run_concurrent_test()) - - # Assert - assert len(results) == 5 - for i, result in enumerate(results): - assert result == [float(i)] * 5 - - def test_processor_add_args_functionality(self): - """Test add_args static method functionality""" - # Arrange - from trustgraph.embeddings.ollama.processor import Processor - mock_parser = Mock() - - # Act - Processor.add_args(mock_parser) - - # Assert - # Verify model argument was added - model_calls = [call for call in mock_parser.add_argument.call_args_list - if call[0][0] in ['-m', '--model']] - assert len(model_calls) == 1 - model_call = model_calls[0] - assert model_call[1]['default'] == "mxbai-embed-large" - - # Verify ollama argument was added - ollama_calls = [call for call in mock_parser.add_argument.call_args_list - if call[0][0] in ['-r', '--ollama']] - assert len(ollama_calls) == 1 - ollama_call = ollama_calls[0] - assert ollama_call[1]['default'] == 'http://localhost:11434' - - def test_host_url_validation(self): - """Test Ollama host URL validation""" - # Arrange - valid_hosts = [ - "http://localhost:11434", - "http://127.0.0.1:11434", - "http://ollama-server:11434", - "https://secure-ollama.com:443" - ] - - # Act & Assert - for host in valid_hosts: - assert isinstance(host, str) - assert host.startswith(("http://", "https://")) - assert ":" in host # Should have port specification - - def test_embedding_vector_properties(self): - """Test properties of embedding vectors""" - # Arrange - test_embedding = [0.1, -0.2, 0.3, -0.4, 0.5, 0.0, -0.1, 0.8] - - # Act - positive_count = sum(1 for x in test_embedding if x > 0) - negative_count = sum(1 for x in test_embedding if x < 0) - zero_count = sum(1 for x in test_embedding if x == 0) - - # Assert - assert positive_count + negative_count + zero_count == len(test_embedding) - assert all(isinstance(x, (int, float)) for x in test_embedding) - assert all(-1.0 <= x <= 1.0 for x in test_embedding) # Typical embedding range \ No newline at end of file diff --git a/tests/unit/test_embeddings/test_ollama_processor.py b/tests/unit/test_embeddings/test_ollama_processor.py deleted file mode 100644 index 3cf03c91..00000000 --- a/tests/unit/test_embeddings/test_ollama_processor.py +++ /dev/null @@ -1,334 +0,0 @@ -""" -Unit tests for Ollama embeddings processor - -Tests the core business logic of Ollama embedding generation without -relying on external Ollama service infrastructure. -""" - -import pytest -from unittest.mock import patch, Mock, AsyncMock, MagicMock -import os - -from trustgraph.embeddings.ollama.processor import Processor -from trustgraph.schema import EmbeddingsRequest, EmbeddingsResponse, Error -from trustgraph.exceptions import TooManyRequests - - -@pytest.fixture -def mock_ollama_client(): - """Mock Ollama client for testing""" - client = Mock() - client.embed.return_value = Mock( - embeddings=[0.1, 0.2, -0.3, 0.4, -0.5, 0.6, 0.7, -0.8, 0.9, -1.0] - ) - return client - - -@pytest.fixture -def processor_params(): - """Default parameters for Ollama processor""" - return { - "model": "test-embed-model", - "ollama": "http://localhost:11434", - "input_queue": "test-input", - "output_queue": "test-output", - "subscriber": "test-subscriber" - } - - -@pytest.fixture -def mock_message(): - """Mock Pulsar message""" - message = Mock() - message.properties.return_value = {"id": "test-msg-123"} - message.value.return_value = EmbeddingsRequest(text="Test embedding text") - return message - - -class TestOllamaEmbeddingsProcessor: - """Test cases for Ollama embeddings processor""" - - @patch('trustgraph.embeddings.ollama.processor.Client') - def test_processor_initialization_default_params(self, mock_client_class): - """Test processor initialization with default parameters""" - # Arrange - mock_client = Mock() - mock_client_class.return_value = mock_client - - # Act - processor = Processor() - - # Assert - assert processor.model == "mxbai-embed-large" # default model - mock_client_class.assert_called_once_with(host='http://localhost:11434') - assert processor.client == mock_client - - @patch('trustgraph.embeddings.ollama.processor.Client') - def test_processor_initialization_custom_params(self, mock_client_class, processor_params): - """Test processor initialization with custom parameters""" - # Arrange - mock_client = Mock() - mock_client_class.return_value = mock_client - - # Act - processor = Processor(**processor_params) - - # Assert - assert processor.model == "test-embed-model" - mock_client_class.assert_called_once_with(host="http://localhost:11434") - assert processor.client == mock_client - - @patch('trustgraph.embeddings.ollama.processor.Client') - def test_processor_initialization_env_variable(self, mock_client_class): - """Test processor uses OLLAMA_HOST environment variable""" - # Arrange - mock_client = Mock() - mock_client_class.return_value = mock_client - - with patch.dict(os.environ, {'OLLAMA_HOST': 'http://custom-host:8080'}): - # Act - processor = Processor() - - # Assert - mock_client_class.assert_called_once_with(host='http://custom-host:8080') - - @patch('trustgraph.embeddings.ollama.processor.Client') - async def test_handle_successful_embedding(self, mock_client_class, mock_message): - """Test successful embedding generation""" - # Arrange - mock_client = Mock() - mock_client.embed.return_value = Mock( - embeddings=[0.1, 0.2, -0.3, 0.4, -0.5] - ) - mock_client_class.return_value = mock_client - - processor = Processor(model="test-model") - processor.send = AsyncMock() - - # Act - await processor.handle(mock_message) - - # Assert - mock_client.embed.assert_called_once_with( - model="test-model", - input="Test embedding text" - ) - - # Verify response sent - processor.send.assert_called_once() - call_args = processor.send.call_args - response = call_args[0][0] - properties = call_args[1]["properties"] - - assert isinstance(response, EmbeddingsResponse) - assert response.error is None - assert response.vectors == [0.1, 0.2, -0.3, 0.4, -0.5] - assert properties["id"] == "test-msg-123" - - @patch('trustgraph.embeddings.ollama.processor.Client') - async def test_handle_client_error(self, mock_client_class, mock_message): - """Test handling of Ollama client errors""" - # Arrange - mock_client = Mock() - mock_client.embed.side_effect = Exception("Ollama connection failed") - mock_client_class.return_value = mock_client - - processor = Processor() - processor.send = AsyncMock() - - # Act & Assert - with pytest.raises(Exception, match="Ollama connection failed"): - await processor.handle(mock_message) - - @patch('trustgraph.embeddings.ollama.processor.Client') - async def test_handle_model_not_found(self, mock_client_class, mock_message): - """Test handling when Ollama model is not found""" - # Arrange - mock_client = Mock() - mock_client.embed.side_effect = Exception("Model 'unknown-model' not found") - mock_client_class.return_value = mock_client - - processor = Processor(model="unknown-model") - processor.send = AsyncMock() - - # Act & Assert - with pytest.raises(Exception, match="Model 'unknown-model' not found"): - await processor.handle(mock_message) - - @patch('trustgraph.embeddings.ollama.processor.Client') - async def test_handle_empty_text(self, mock_client_class): - """Test handling of empty text input""" - # Arrange - mock_client = Mock() - mock_client.embed.return_value = Mock(embeddings=[]) - mock_client_class.return_value = mock_client - - processor = Processor() - processor.send = AsyncMock() - - empty_message = Mock() - empty_message.properties.return_value = {"id": "empty-msg"} - empty_message.value.return_value = EmbeddingsRequest(text="") - - # Act - await processor.handle(empty_message) - - # Assert - mock_client.embed.assert_called_once_with( - model="mxbai-embed-large", - input="" - ) - - processor.send.assert_called_once() - response = processor.send.call_args[0][0] - assert response.vectors == [] - - @patch('trustgraph.embeddings.ollama.processor.Client') - async def test_handle_large_text_input(self, mock_client_class): - """Test handling of large text input""" - # Arrange - mock_client = Mock() - mock_client.embed.return_value = Mock( - embeddings=[0.1] * 1536 # Large embedding vector - ) - mock_client_class.return_value = mock_client - - processor = Processor() - processor.send = AsyncMock() - - large_text = "A" * 10000 # Large text input - large_message = Mock() - large_message.properties.return_value = {"id": "large-msg"} - large_message.value.return_value = EmbeddingsRequest(text=large_text) - - # Act - await processor.handle(large_message) - - # Assert - mock_client.embed.assert_called_once_with( - model="mxbai-embed-large", - input=large_text - ) - - processor.send.assert_called_once() - response = processor.send.call_args[0][0] - assert len(response.vectors) == 1536 - - @patch('trustgraph.embeddings.ollama.processor.Client') - def test_embedding_vector_consistency(self, mock_client_class): - """Test that embedding vectors have consistent dimensions""" - # Arrange - mock_client = Mock() - mock_client_class.return_value = mock_client - - processor = Processor(model="test-model") - - # Assert - assert processor.model == "test-model" - assert processor.client == mock_client - - @patch('trustgraph.embeddings.ollama.processor.Client') - async def test_handle_unicode_text(self, mock_client_class): - """Test handling of Unicode text input""" - # Arrange - mock_client = Mock() - mock_client.embed.return_value = Mock( - embeddings=[0.2, -0.1, 0.3, -0.4, 0.5] - ) - mock_client_class.return_value = mock_client - - processor = Processor() - processor.send = AsyncMock() - - unicode_text = "Hello 世界! 🌍 Émoji test" - unicode_message = Mock() - unicode_message.properties.return_value = {"id": "unicode-msg"} - unicode_message.value.return_value = EmbeddingsRequest(text=unicode_text) - - # Act - await processor.handle(unicode_message) - - # Assert - mock_client.embed.assert_called_once_with( - model="mxbai-embed-large", - input=unicode_text - ) - - processor.send.assert_called_once() - response = processor.send.call_args[0][0] - assert response.vectors == [0.2, -0.1, 0.3, -0.4, 0.5] - - def test_add_args_method(self): - """Test that add_args method adds correct arguments""" - # Arrange - mock_parser = Mock() - - # Act - Processor.add_args(mock_parser) - - # Assert - # Verify model argument added - model_calls = [call for call in mock_parser.add_argument.call_args_list - if call[0][0] in ['-m', '--model']] - assert len(model_calls) == 1 - model_call = model_calls[0] - assert model_call[1]['default'] == "mxbai-embed-large" - - # Verify ollama argument added - ollama_calls = [call for call in mock_parser.add_argument.call_args_list - if call[0][0] in ['-r', '--ollama']] - assert len(ollama_calls) == 1 - ollama_call = ollama_calls[0] - assert ollama_call[1]['default'] == 'http://localhost:11434' - - @patch('trustgraph.embeddings.ollama.processor.Client') - async def test_message_properties_handling(self, mock_client_class, mock_message): - """Test proper handling of message properties""" - # Arrange - mock_client = Mock() - mock_client.embed.return_value = Mock(embeddings=[0.1, 0.2, 0.3]) - mock_client_class.return_value = mock_client - - processor = Processor() - processor.send = AsyncMock() - - # Act - await processor.handle(mock_message) - - # Assert - mock_message.properties.assert_called_once() - - # Verify properties passed to send - call_properties = processor.send.call_args[1]["properties"] - assert call_properties["id"] == "test-msg-123" - - @patch('trustgraph.embeddings.ollama.processor.Client') - async def test_concurrent_requests_handling(self, mock_client_class): - """Test that processor can handle multiple concurrent requests""" - # Arrange - mock_client = Mock() - mock_client.embed.return_value = Mock(embeddings=[0.1, 0.2, 0.3]) - mock_client_class.return_value = mock_client - - processor = Processor() - processor.send = AsyncMock() - - # Create multiple messages - messages = [] - for i in range(3): - msg = Mock() - msg.properties.return_value = {"id": f"msg-{i}"} - msg.value.return_value = EmbeddingsRequest(text=f"Text {i}") - messages.append(msg) - - # Act - Handle messages concurrently - import asyncio - await asyncio.gather(*[processor.handle(msg) for msg in messages]) - - # Assert - assert mock_client.embed.call_count == 3 - assert processor.send.call_count == 3 - - # Verify each message was handled correctly - for i, call in enumerate(mock_client.embed.call_args_list): - assert call[1]["input"] == f"Text {i}" \ No newline at end of file diff --git a/trustgraph-flow/trustgraph/embeddings/ollama/processor.py b/trustgraph-flow/trustgraph/embeddings/ollama/processor.py index 86787316..3c0776f9 100755 --- a/trustgraph-flow/trustgraph/embeddings/ollama/processor.py +++ b/trustgraph-flow/trustgraph/embeddings/ollama/processor.py @@ -3,81 +3,46 @@ Embeddings service, applies an embeddings model hosted on a local Ollama. Input is text, output is embeddings vector. """ +from ... base import EmbeddingsService -from ... schema import EmbeddingsRequest, EmbeddingsResponse -from ... schema import embeddings_request_queue, embeddings_response_queue -from ... log_level import LogLevel -from ... base import ConsumerProducer from ollama import Client import os -module = "embeddings" +default_ident = "embeddings" -default_input_queue = embeddings_request_queue -default_output_queue = embeddings_response_queue -default_subscriber = module default_model="mxbai-embed-large" default_ollama = os.getenv("OLLAMA_HOST", 'http://localhost:11434') -class Processor(ConsumerProducer): +class Processor(EmbeddingsService): def __init__(self, **params): - input_queue = params.get("input_queue", default_input_queue) - output_queue = params.get("output_queue", default_output_queue) - subscriber = params.get("subscriber", default_subscriber) - - ollama = params.get("ollama", default_ollama) model = params.get("model", default_model) + ollama = params.get("ollama", default_ollama) super(Processor, self).__init__( **params | { - "input_queue": input_queue, - "output_queue": output_queue, - "subscriber": subscriber, - "input_schema": EmbeddingsRequest, - "output_schema": EmbeddingsResponse, "ollama": ollama, - "model": model, + "model": model } ) self.client = Client(host=ollama) self.model = model - async def handle(self, msg): + async def on_embeddings(self, text): - v = msg.value() - - # Sender-produced ID - - id = msg.properties()["id"] - - print(f"Handling input {id}...", flush=True) - - text = v.text embeds = self.client.embed( model = self.model, input = text ) - print("Send response...", flush=True) - r = EmbeddingsResponse( - vectors=embeds.embeddings, - error=None, - ) - - await self.send(r, properties={"id": id}) - - print("Done.", flush=True) + return embeds.embeddings @staticmethod def add_args(parser): - ConsumerProducer.add_args( - parser, default_input_queue, default_subscriber, - default_output_queue, - ) + EmbeddingsService.add_args(parser) parser.add_argument( '-m', '--model', @@ -93,5 +58,6 @@ class Processor(ConsumerProducer): def run(): - Processor.launch(module, __doc__) + Processor.launch(default_ident, __doc__) +