From 76fd0534fcb17bd8a4079c07bd5b04beec889a7b Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 14 Jul 2025 22:41:50 +0100 Subject: [PATCH] Fixing storage and adding tests --- .../test_doc_embeddings_pinecone_query.py | 558 ++++++++++++++++++ .../test_doc_embeddings_pinecone_storage.py | 536 +++++++++++++++++ .../query/doc_embeddings/pinecone/service.py | 82 +-- .../storage/doc_embeddings/pinecone/write.py | 148 +++-- 4 files changed, 1183 insertions(+), 141 deletions(-) create mode 100644 tests/unit/test_query/test_doc_embeddings_pinecone_query.py create mode 100644 tests/unit/test_storage/test_doc_embeddings_pinecone_storage.py diff --git a/tests/unit/test_query/test_doc_embeddings_pinecone_query.py b/tests/unit/test_query/test_doc_embeddings_pinecone_query.py new file mode 100644 index 00000000..92551587 --- /dev/null +++ b/tests/unit/test_query/test_doc_embeddings_pinecone_query.py @@ -0,0 +1,558 @@ +""" +Tests for Pinecone document embeddings query service +""" + +import pytest +from unittest.mock import MagicMock, patch + +from trustgraph.query.doc_embeddings.pinecone.service import Processor + + +class TestPineconeDocEmbeddingsQueryProcessor: + """Test cases for Pinecone document embeddings query processor""" + + @pytest.fixture + def mock_query_message(self): + """Create a mock query message for testing""" + message = MagicMock() + message.vectors = [ + [0.1, 0.2, 0.3], + [0.4, 0.5, 0.6] + ] + message.limit = 5 + message.user = 'test_user' + message.collection = 'test_collection' + return message + + @pytest.fixture + def processor(self): + """Create a processor instance for testing""" + with patch('trustgraph.query.doc_embeddings.pinecone.service.Pinecone') as mock_pinecone_class: + mock_pinecone = MagicMock() + mock_pinecone_class.return_value = mock_pinecone + + processor = Processor( + taskgroup=MagicMock(), + id='test-pinecone-de-query', + api_key='test-api-key' + ) + + return processor + + @patch('trustgraph.query.doc_embeddings.pinecone.service.Pinecone') + @patch('trustgraph.query.doc_embeddings.pinecone.service.default_api_key', 'env-api-key') + def test_processor_initialization_with_defaults(self, mock_pinecone_class): + """Test processor initialization with default parameters""" + mock_pinecone = MagicMock() + mock_pinecone_class.return_value = mock_pinecone + taskgroup_mock = MagicMock() + + processor = Processor(taskgroup=taskgroup_mock) + + mock_pinecone_class.assert_called_once_with(api_key='env-api-key') + assert processor.pinecone == mock_pinecone + assert processor.api_key == 'env-api-key' + + @patch('trustgraph.query.doc_embeddings.pinecone.service.Pinecone') + def test_processor_initialization_with_custom_params(self, mock_pinecone_class): + """Test processor initialization with custom parameters""" + mock_pinecone = MagicMock() + mock_pinecone_class.return_value = mock_pinecone + taskgroup_mock = MagicMock() + + processor = Processor( + taskgroup=taskgroup_mock, + api_key='custom-api-key' + ) + + mock_pinecone_class.assert_called_once_with(api_key='custom-api-key') + assert processor.api_key == 'custom-api-key' + + @patch('trustgraph.query.doc_embeddings.pinecone.service.PineconeGRPC') + def test_processor_initialization_with_url(self, mock_pinecone_grpc_class): + """Test processor initialization with custom URL (GRPC mode)""" + mock_pinecone = MagicMock() + mock_pinecone_grpc_class.return_value = mock_pinecone + taskgroup_mock = MagicMock() + + processor = Processor( + taskgroup=taskgroup_mock, + api_key='test-api-key', + url='https://custom-host.pinecone.io' + ) + + mock_pinecone_grpc_class.assert_called_once_with( + api_key='test-api-key', + host='https://custom-host.pinecone.io' + ) + assert processor.pinecone == mock_pinecone + assert processor.url == 'https://custom-host.pinecone.io' + + @patch('trustgraph.query.doc_embeddings.pinecone.service.default_api_key', 'not-specified') + def test_processor_initialization_missing_api_key(self): + """Test processor initialization fails with missing API key""" + taskgroup_mock = MagicMock() + + with pytest.raises(RuntimeError, match="Pinecone API key must be specified"): + Processor(taskgroup=taskgroup_mock) + + @pytest.mark.asyncio + async def test_query_document_embeddings_single_vector(self, processor): + """Test querying document embeddings with a single vector""" + message = MagicMock() + message.vectors = [[0.1, 0.2, 0.3]] + message.limit = 3 + message.user = 'test_user' + message.collection = 'test_collection' + + # Mock index and query results + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + mock_results = MagicMock() + mock_results.matches = [ + MagicMock(metadata={'doc': 'First document chunk'}), + MagicMock(metadata={'doc': 'Second document chunk'}), + MagicMock(metadata={'doc': 'Third document chunk'}) + ] + mock_index.query.return_value = mock_results + + chunks = await processor.query_document_embeddings(message) + + # Verify index was accessed correctly + expected_index_name = "d-test_user-test_collection-3" + processor.pinecone.Index.assert_called_once_with(expected_index_name) + + # Verify query parameters + mock_index.query.assert_called_once_with( + vector=[0.1, 0.2, 0.3], + top_k=3, + include_values=False, + include_metadata=True + ) + + # Verify results + assert len(chunks) == 3 + assert chunks[0] == 'First document chunk' + assert chunks[1] == 'Second document chunk' + assert chunks[2] == 'Third document chunk' + + @pytest.mark.asyncio + async def test_query_document_embeddings_multiple_vectors(self, processor, mock_query_message): + """Test querying document embeddings with multiple vectors""" + # Mock index and query results + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + # First query results + mock_results1 = MagicMock() + mock_results1.matches = [ + MagicMock(metadata={'doc': 'Document chunk 1'}), + MagicMock(metadata={'doc': 'Document chunk 2'}) + ] + + # Second query results + mock_results2 = MagicMock() + mock_results2.matches = [ + MagicMock(metadata={'doc': 'Document chunk 3'}), + MagicMock(metadata={'doc': 'Document chunk 4'}) + ] + + mock_index.query.side_effect = [mock_results1, mock_results2] + + chunks = await processor.query_document_embeddings(mock_query_message) + + # Verify both queries were made + assert mock_index.query.call_count == 2 + + # Verify results from both queries + assert len(chunks) == 4 + assert 'Document chunk 1' in chunks + assert 'Document chunk 2' in chunks + assert 'Document chunk 3' in chunks + assert 'Document chunk 4' in chunks + + @pytest.mark.asyncio + async def test_query_document_embeddings_limit_handling(self, processor): + """Test that query respects the limit parameter""" + message = MagicMock() + message.vectors = [[0.1, 0.2, 0.3]] + message.limit = 2 + message.user = 'test_user' + message.collection = 'test_collection' + + # Mock index with many results + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + mock_results = MagicMock() + mock_results.matches = [ + MagicMock(metadata={'doc': f'Document chunk {i}'}) for i in range(10) + ] + mock_index.query.return_value = mock_results + + chunks = await processor.query_document_embeddings(message) + + # Verify limit is passed to query + mock_index.query.assert_called_once() + call_args = mock_index.query.call_args + assert call_args[1]['top_k'] == 2 + + # Results should contain all returned chunks (limit is applied by Pinecone) + assert len(chunks) == 10 + + @pytest.mark.asyncio + async def test_query_document_embeddings_zero_limit(self, processor): + """Test querying with zero limit returns empty results""" + message = MagicMock() + message.vectors = [[0.1, 0.2, 0.3]] + message.limit = 0 + message.user = 'test_user' + message.collection = 'test_collection' + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + chunks = await processor.query_document_embeddings(message) + + # Verify no query was made and empty result returned + mock_index.query.assert_not_called() + assert chunks == [] + + @pytest.mark.asyncio + async def test_query_document_embeddings_negative_limit(self, processor): + """Test querying with negative limit returns empty results""" + message = MagicMock() + message.vectors = [[0.1, 0.2, 0.3]] + message.limit = -1 + message.user = 'test_user' + message.collection = 'test_collection' + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + chunks = await processor.query_document_embeddings(message) + + # Verify no query was made and empty result returned + mock_index.query.assert_not_called() + assert chunks == [] + + @pytest.mark.asyncio + async def test_query_document_embeddings_different_vector_dimensions(self, processor): + """Test querying with vectors of different dimensions""" + message = MagicMock() + message.vectors = [ + [0.1, 0.2], # 2D vector + [0.3, 0.4, 0.5, 0.6] # 4D vector + ] + message.limit = 5 + message.user = 'test_user' + message.collection = 'test_collection' + + mock_index_2d = MagicMock() + mock_index_4d = MagicMock() + + def mock_index_side_effect(name): + if name.endswith("-2"): + return mock_index_2d + elif name.endswith("-4"): + return mock_index_4d + + processor.pinecone.Index.side_effect = mock_index_side_effect + + # Mock results for different dimensions + mock_results_2d = MagicMock() + mock_results_2d.matches = [MagicMock(metadata={'doc': 'Document from 2D index'})] + mock_index_2d.query.return_value = mock_results_2d + + mock_results_4d = MagicMock() + mock_results_4d.matches = [MagicMock(metadata={'doc': 'Document from 4D index'})] + mock_index_4d.query.return_value = mock_results_4d + + chunks = await processor.query_document_embeddings(message) + + # Verify different indexes were used + assert processor.pinecone.Index.call_count == 2 + mock_index_2d.query.assert_called_once() + mock_index_4d.query.assert_called_once() + + # Verify results from both dimensions + assert 'Document from 2D index' in chunks + assert 'Document from 4D index' in chunks + + @pytest.mark.asyncio + async def test_query_document_embeddings_empty_vectors_list(self, processor): + """Test querying with empty vectors list""" + message = MagicMock() + message.vectors = [] + message.limit = 5 + message.user = 'test_user' + message.collection = 'test_collection' + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + chunks = await processor.query_document_embeddings(message) + + # Verify no queries were made and empty result returned + processor.pinecone.Index.assert_not_called() + mock_index.query.assert_not_called() + assert chunks == [] + + @pytest.mark.asyncio + async def test_query_document_embeddings_no_results(self, processor): + """Test querying when index returns no results""" + message = MagicMock() + message.vectors = [[0.1, 0.2, 0.3]] + message.limit = 5 + message.user = 'test_user' + message.collection = 'test_collection' + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + mock_results = MagicMock() + mock_results.matches = [] + mock_index.query.return_value = mock_results + + chunks = await processor.query_document_embeddings(message) + + # Verify empty results + assert chunks == [] + + @pytest.mark.asyncio + async def test_query_document_embeddings_unicode_content(self, processor): + """Test querying document embeddings with Unicode content results""" + message = MagicMock() + message.vectors = [[0.1, 0.2, 0.3]] + message.limit = 2 + message.user = 'test_user' + message.collection = 'test_collection' + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + mock_results = MagicMock() + mock_results.matches = [ + MagicMock(metadata={'doc': 'Document with Unicode: éñ中文🚀'}), + MagicMock(metadata={'doc': 'Regular ASCII document'}) + ] + mock_index.query.return_value = mock_results + + chunks = await processor.query_document_embeddings(message) + + # Verify Unicode content is properly handled + assert len(chunks) == 2 + assert 'Document with Unicode: éñ中文🚀' in chunks + assert 'Regular ASCII document' in chunks + + @pytest.mark.asyncio + async def test_query_document_embeddings_large_content(self, processor): + """Test querying document embeddings with large content results""" + message = MagicMock() + message.vectors = [[0.1, 0.2, 0.3]] + message.limit = 1 + message.user = 'test_user' + message.collection = 'test_collection' + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + # Create a large document content + large_content = "A" * 10000 # 10KB of content + mock_results = MagicMock() + mock_results.matches = [ + MagicMock(metadata={'doc': large_content}) + ] + mock_index.query.return_value = mock_results + + chunks = await processor.query_document_embeddings(message) + + # Verify large content is properly handled + assert len(chunks) == 1 + assert chunks[0] == large_content + + @pytest.mark.asyncio + async def test_query_document_embeddings_mixed_content_types(self, processor): + """Test querying document embeddings with mixed content types""" + message = MagicMock() + message.vectors = [[0.1, 0.2, 0.3]] + message.limit = 5 + message.user = 'test_user' + message.collection = 'test_collection' + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + mock_results = MagicMock() + mock_results.matches = [ + MagicMock(metadata={'doc': 'Short text'}), + MagicMock(metadata={'doc': 'A' * 1000}), # Long text + MagicMock(metadata={'doc': 'Text with numbers: 123 and symbols: @#$'}), + MagicMock(metadata={'doc': ' Whitespace text '}), + MagicMock(metadata={'doc': ''}) # Empty string + ] + mock_index.query.return_value = mock_results + + chunks = await processor.query_document_embeddings(message) + + # Verify all content types are properly handled + assert len(chunks) == 5 + assert 'Short text' in chunks + assert 'A' * 1000 in chunks + assert 'Text with numbers: 123 and symbols: @#$' in chunks + assert ' Whitespace text ' in chunks + assert '' in chunks + + @pytest.mark.asyncio + async def test_query_document_embeddings_exception_handling(self, processor): + """Test that exceptions are properly raised""" + message = MagicMock() + message.vectors = [[0.1, 0.2, 0.3]] + message.limit = 5 + message.user = 'test_user' + message.collection = 'test_collection' + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + mock_index.query.side_effect = Exception("Query failed") + + with pytest.raises(Exception, match="Query failed"): + await processor.query_document_embeddings(message) + + @pytest.mark.asyncio + async def test_query_document_embeddings_index_access_failure(self, processor): + """Test handling of index access failure""" + message = MagicMock() + message.vectors = [[0.1, 0.2, 0.3]] + message.limit = 5 + message.user = 'test_user' + message.collection = 'test_collection' + + processor.pinecone.Index.side_effect = Exception("Index access failed") + + with pytest.raises(Exception, match="Index access failed"): + await processor.query_document_embeddings(message) + + @pytest.mark.asyncio + async def test_query_document_embeddings_vector_accumulation(self, processor): + """Test that results from multiple vectors are properly accumulated""" + message = MagicMock() + message.vectors = [ + [0.1, 0.2, 0.3], + [0.4, 0.5, 0.6], + [0.7, 0.8, 0.9] + ] + message.limit = 2 + message.user = 'test_user' + message.collection = 'test_collection' + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + # Each query returns different results + mock_results1 = MagicMock() + mock_results1.matches = [ + MagicMock(metadata={'doc': 'Doc from vector 1.1'}), + MagicMock(metadata={'doc': 'Doc from vector 1.2'}) + ] + + mock_results2 = MagicMock() + mock_results2.matches = [ + MagicMock(metadata={'doc': 'Doc from vector 2.1'}) + ] + + mock_results3 = MagicMock() + mock_results3.matches = [ + MagicMock(metadata={'doc': 'Doc from vector 3.1'}), + MagicMock(metadata={'doc': 'Doc from vector 3.2'}), + MagicMock(metadata={'doc': 'Doc from vector 3.3'}) + ] + + mock_index.query.side_effect = [mock_results1, mock_results2, mock_results3] + + chunks = await processor.query_document_embeddings(message) + + # Verify all queries were made + assert mock_index.query.call_count == 3 + + # Verify all results are accumulated + assert len(chunks) == 6 + assert 'Doc from vector 1.1' in chunks + assert 'Doc from vector 1.2' in chunks + assert 'Doc from vector 2.1' in chunks + assert 'Doc from vector 3.1' in chunks + assert 'Doc from vector 3.2' in chunks + assert 'Doc from vector 3.3' in chunks + + def test_add_args_method(self): + """Test that add_args properly configures argument parser""" + from argparse import ArgumentParser + from unittest.mock import patch + + parser = ArgumentParser() + + # Mock the parent class add_args method + with patch('trustgraph.query.doc_embeddings.pinecone.service.DocumentEmbeddingsQueryService.add_args') as mock_parent_add_args: + Processor.add_args(parser) + + # Verify parent add_args was called + mock_parent_add_args.assert_called_once() + + # Verify our specific arguments were added + args = parser.parse_args([]) + + assert hasattr(args, 'api_key') + assert args.api_key == 'not-specified' # Default value when no env var + assert hasattr(args, 'url') + assert args.url is None + + def test_add_args_with_custom_values(self): + """Test add_args with custom command line values""" + from argparse import ArgumentParser + from unittest.mock import patch + + parser = ArgumentParser() + + with patch('trustgraph.query.doc_embeddings.pinecone.service.DocumentEmbeddingsQueryService.add_args'): + Processor.add_args(parser) + + # Test parsing with custom values + args = parser.parse_args([ + '--api-key', 'custom-api-key', + '--url', 'https://custom-host.pinecone.io' + ]) + + assert args.api_key == 'custom-api-key' + assert args.url == 'https://custom-host.pinecone.io' + + def test_add_args_short_form(self): + """Test add_args with short form arguments""" + from argparse import ArgumentParser + from unittest.mock import patch + + parser = ArgumentParser() + + with patch('trustgraph.query.doc_embeddings.pinecone.service.DocumentEmbeddingsQueryService.add_args'): + Processor.add_args(parser) + + # Test parsing with short form + args = parser.parse_args([ + '-a', 'short-api-key', + '-u', 'https://short-host.pinecone.io' + ]) + + assert args.api_key == 'short-api-key' + assert args.url == 'https://short-host.pinecone.io' + + @patch('trustgraph.query.doc_embeddings.pinecone.service.Processor.launch') + def test_run_function(self, mock_launch): + """Test the run function calls Processor.launch with correct parameters""" + from trustgraph.query.doc_embeddings.pinecone.service import run, default_ident + + run() + + mock_launch.assert_called_once_with( + default_ident, + "\nDocument embeddings query service. Input is vector, output is an array\nof chunks. Pinecone implementation.\n" + ) \ No newline at end of file diff --git a/tests/unit/test_storage/test_doc_embeddings_pinecone_storage.py b/tests/unit/test_storage/test_doc_embeddings_pinecone_storage.py new file mode 100644 index 00000000..6c4ddb6b --- /dev/null +++ b/tests/unit/test_storage/test_doc_embeddings_pinecone_storage.py @@ -0,0 +1,536 @@ +""" +Tests for Pinecone document embeddings storage service +""" + +import pytest +from unittest.mock import MagicMock, patch +import uuid + +from trustgraph.storage.doc_embeddings.pinecone.write import Processor +from trustgraph.schema import ChunkEmbeddings + + +class TestPineconeDocEmbeddingsStorageProcessor: + """Test cases for Pinecone document embeddings storage processor""" + + @pytest.fixture + def mock_message(self): + """Create a mock message for testing""" + message = MagicMock() + message.metadata = MagicMock() + message.metadata.user = 'test_user' + message.metadata.collection = 'test_collection' + + # Create test document embeddings + chunk1 = ChunkEmbeddings( + chunk=b"This is the first document chunk", + vectors=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + ) + chunk2 = ChunkEmbeddings( + chunk=b"This is the second document chunk", + vectors=[[0.7, 0.8, 0.9]] + ) + message.chunks = [chunk1, chunk2] + + return message + + @pytest.fixture + def processor(self): + """Create a processor instance for testing""" + with patch('trustgraph.storage.doc_embeddings.pinecone.write.Pinecone') as mock_pinecone_class: + mock_pinecone = MagicMock() + mock_pinecone_class.return_value = mock_pinecone + + processor = Processor( + taskgroup=MagicMock(), + id='test-pinecone-de-storage', + api_key='test-api-key' + ) + + return processor + + @patch('trustgraph.storage.doc_embeddings.pinecone.write.Pinecone') + @patch('trustgraph.storage.doc_embeddings.pinecone.write.default_api_key', 'env-api-key') + def test_processor_initialization_with_defaults(self, mock_pinecone_class): + """Test processor initialization with default parameters""" + mock_pinecone = MagicMock() + mock_pinecone_class.return_value = mock_pinecone + taskgroup_mock = MagicMock() + + processor = Processor(taskgroup=taskgroup_mock) + + mock_pinecone_class.assert_called_once_with(api_key='env-api-key') + assert processor.pinecone == mock_pinecone + assert processor.api_key == 'env-api-key' + assert processor.cloud == 'aws' + assert processor.region == 'us-east-1' + + @patch('trustgraph.storage.doc_embeddings.pinecone.write.Pinecone') + def test_processor_initialization_with_custom_params(self, mock_pinecone_class): + """Test processor initialization with custom parameters""" + mock_pinecone = MagicMock() + mock_pinecone_class.return_value = mock_pinecone + taskgroup_mock = MagicMock() + + processor = Processor( + taskgroup=taskgroup_mock, + api_key='custom-api-key', + cloud='gcp', + region='us-west1' + ) + + mock_pinecone_class.assert_called_once_with(api_key='custom-api-key') + assert processor.api_key == 'custom-api-key' + assert processor.cloud == 'gcp' + assert processor.region == 'us-west1' + + @patch('trustgraph.storage.doc_embeddings.pinecone.write.PineconeGRPC') + def test_processor_initialization_with_url(self, mock_pinecone_grpc_class): + """Test processor initialization with custom URL (GRPC mode)""" + mock_pinecone = MagicMock() + mock_pinecone_grpc_class.return_value = mock_pinecone + taskgroup_mock = MagicMock() + + processor = Processor( + taskgroup=taskgroup_mock, + api_key='test-api-key', + url='https://custom-host.pinecone.io' + ) + + mock_pinecone_grpc_class.assert_called_once_with( + api_key='test-api-key', + host='https://custom-host.pinecone.io' + ) + assert processor.pinecone == mock_pinecone + assert processor.url == 'https://custom-host.pinecone.io' + + @patch('trustgraph.storage.doc_embeddings.pinecone.write.default_api_key', 'not-specified') + def test_processor_initialization_missing_api_key(self): + """Test processor initialization fails with missing API key""" + taskgroup_mock = MagicMock() + + with pytest.raises(RuntimeError, match="Pinecone API key must be specified"): + Processor(taskgroup=taskgroup_mock) + + @pytest.mark.asyncio + async def test_store_document_embeddings_single_chunk(self, processor): + """Test storing document embeddings for a single chunk""" + message = MagicMock() + message.metadata = MagicMock() + message.metadata.user = 'test_user' + message.metadata.collection = 'test_collection' + + chunk = ChunkEmbeddings( + chunk=b"Test document content", + vectors=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + ) + message.chunks = [chunk] + + # Mock index operations + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + processor.pinecone.has_index.return_value = True + + with patch('uuid.uuid4', side_effect=['id1', 'id2']): + await processor.store_document_embeddings(message) + + # Verify index name and operations + expected_index_name = "d-test_user-test_collection-3" + processor.pinecone.Index.assert_called_with(expected_index_name) + + # Verify upsert was called for each vector + assert mock_index.upsert.call_count == 2 + + # Check first vector upsert + first_call = mock_index.upsert.call_args_list[0] + first_vectors = first_call[1]['vectors'] + assert len(first_vectors) == 1 + assert first_vectors[0]['id'] == 'id1' + assert first_vectors[0]['values'] == [0.1, 0.2, 0.3] + assert first_vectors[0]['metadata']['doc'] == "Test document content" + + # Check second vector upsert + second_call = mock_index.upsert.call_args_list[1] + second_vectors = second_call[1]['vectors'] + assert len(second_vectors) == 1 + assert second_vectors[0]['id'] == 'id2' + assert second_vectors[0]['values'] == [0.4, 0.5, 0.6] + assert second_vectors[0]['metadata']['doc'] == "Test document content" + + @pytest.mark.asyncio + async def test_store_document_embeddings_multiple_chunks(self, processor, mock_message): + """Test storing document embeddings for multiple chunks""" + # Mock index operations + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + processor.pinecone.has_index.return_value = True + + with patch('uuid.uuid4', side_effect=['id1', 'id2', 'id3']): + await processor.store_document_embeddings(mock_message) + + # Verify upsert was called for each vector (3 total) + assert mock_index.upsert.call_count == 3 + + # Verify document content in metadata + calls = mock_index.upsert.call_args_list + assert calls[0][1]['vectors'][0]['metadata']['doc'] == "This is the first document chunk" + assert calls[1][1]['vectors'][0]['metadata']['doc'] == "This is the first document chunk" + assert calls[2][1]['vectors'][0]['metadata']['doc'] == "This is the second document chunk" + + @pytest.mark.asyncio + async def test_store_document_embeddings_index_creation(self, processor): + """Test automatic index creation when index doesn't exist""" + message = MagicMock() + message.metadata = MagicMock() + message.metadata.user = 'test_user' + message.metadata.collection = 'test_collection' + + chunk = ChunkEmbeddings( + chunk=b"Test document content", + vectors=[[0.1, 0.2, 0.3]] + ) + message.chunks = [chunk] + + # Mock index doesn't exist initially + processor.pinecone.has_index.return_value = False + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + # Mock index creation + processor.pinecone.describe_index.return_value.status = {"ready": True} + + with patch('uuid.uuid4', return_value='test-id'): + await processor.store_document_embeddings(message) + + # Verify index creation was called + expected_index_name = "d-test_user-test_collection-3" + processor.pinecone.create_index.assert_called_once() + create_call = processor.pinecone.create_index.call_args + assert create_call[1]['name'] == expected_index_name + assert create_call[1]['dimension'] == 3 + assert create_call[1]['metric'] == "cosine" + + @pytest.mark.asyncio + async def test_store_document_embeddings_empty_chunk(self, processor): + """Test storing document embeddings with empty chunk (should be skipped)""" + message = MagicMock() + message.metadata = MagicMock() + message.metadata.user = 'test_user' + message.metadata.collection = 'test_collection' + + chunk = ChunkEmbeddings( + chunk=b"", + vectors=[[0.1, 0.2, 0.3]] + ) + message.chunks = [chunk] + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + await processor.store_document_embeddings(message) + + # Verify no upsert was called for empty chunk + mock_index.upsert.assert_not_called() + + @pytest.mark.asyncio + async def test_store_document_embeddings_none_chunk(self, processor): + """Test storing document embeddings with None chunk (should be skipped)""" + message = MagicMock() + message.metadata = MagicMock() + message.metadata.user = 'test_user' + message.metadata.collection = 'test_collection' + + chunk = ChunkEmbeddings( + chunk=None, + vectors=[[0.1, 0.2, 0.3]] + ) + message.chunks = [chunk] + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + await processor.store_document_embeddings(message) + + # Verify no upsert was called for None chunk + mock_index.upsert.assert_not_called() + + @pytest.mark.asyncio + async def test_store_document_embeddings_empty_decoded_chunk(self, processor): + """Test storing document embeddings with chunk that decodes to empty string""" + message = MagicMock() + message.metadata = MagicMock() + message.metadata.user = 'test_user' + message.metadata.collection = 'test_collection' + + chunk = ChunkEmbeddings( + chunk=b"", # Empty bytes + vectors=[[0.1, 0.2, 0.3]] + ) + message.chunks = [chunk] + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + await processor.store_document_embeddings(message) + + # Verify no upsert was called for empty decoded chunk + mock_index.upsert.assert_not_called() + + @pytest.mark.asyncio + async def test_store_document_embeddings_different_vector_dimensions(self, processor): + """Test storing document embeddings with different vector dimensions""" + message = MagicMock() + message.metadata = MagicMock() + message.metadata.user = 'test_user' + message.metadata.collection = 'test_collection' + + chunk = ChunkEmbeddings( + chunk=b"Document with mixed dimensions", + vectors=[ + [0.1, 0.2], # 2D vector + [0.3, 0.4, 0.5, 0.6], # 4D vector + [0.7, 0.8, 0.9] # 3D vector + ] + ) + message.chunks = [chunk] + + mock_index_2d = MagicMock() + mock_index_4d = MagicMock() + mock_index_3d = MagicMock() + + def mock_index_side_effect(name): + if name.endswith("-2"): + return mock_index_2d + elif name.endswith("-4"): + return mock_index_4d + elif name.endswith("-3"): + return mock_index_3d + + processor.pinecone.Index.side_effect = mock_index_side_effect + processor.pinecone.has_index.return_value = True + + with patch('uuid.uuid4', side_effect=['id1', 'id2', 'id3']): + await processor.store_document_embeddings(message) + + # Verify different indexes were used for different dimensions + assert processor.pinecone.Index.call_count == 3 + mock_index_2d.upsert.assert_called_once() + mock_index_4d.upsert.assert_called_once() + mock_index_3d.upsert.assert_called_once() + + @pytest.mark.asyncio + async def test_store_document_embeddings_empty_chunks_list(self, processor): + """Test storing document embeddings with empty chunks list""" + message = MagicMock() + message.metadata = MagicMock() + message.metadata.user = 'test_user' + message.metadata.collection = 'test_collection' + message.chunks = [] + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + await processor.store_document_embeddings(message) + + # Verify no operations were performed + processor.pinecone.Index.assert_not_called() + mock_index.upsert.assert_not_called() + + @pytest.mark.asyncio + async def test_store_document_embeddings_chunk_with_no_vectors(self, processor): + """Test storing document embeddings for chunk with no vectors""" + message = MagicMock() + message.metadata = MagicMock() + message.metadata.user = 'test_user' + message.metadata.collection = 'test_collection' + + chunk = ChunkEmbeddings( + chunk=b"Document with no vectors", + vectors=[] + ) + message.chunks = [chunk] + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + + await processor.store_document_embeddings(message) + + # Verify no upsert was called (no vectors to insert) + mock_index.upsert.assert_not_called() + + @pytest.mark.asyncio + async def test_store_document_embeddings_index_creation_failure(self, processor): + """Test handling of index creation failure""" + message = MagicMock() + message.metadata = MagicMock() + message.metadata.user = 'test_user' + message.metadata.collection = 'test_collection' + + chunk = ChunkEmbeddings( + chunk=b"Test document content", + vectors=[[0.1, 0.2, 0.3]] + ) + message.chunks = [chunk] + + # Mock index doesn't exist and creation fails + processor.pinecone.has_index.return_value = False + processor.pinecone.create_index.side_effect = Exception("Index creation failed") + + with pytest.raises(Exception, match="Index creation failed"): + await processor.store_document_embeddings(message) + + @pytest.mark.asyncio + async def test_store_document_embeddings_index_creation_timeout(self, processor): + """Test handling of index creation timeout""" + message = MagicMock() + message.metadata = MagicMock() + message.metadata.user = 'test_user' + message.metadata.collection = 'test_collection' + + chunk = ChunkEmbeddings( + chunk=b"Test document content", + vectors=[[0.1, 0.2, 0.3]] + ) + message.chunks = [chunk] + + # Mock index doesn't exist and never becomes ready + processor.pinecone.has_index.return_value = False + processor.pinecone.describe_index.return_value.status = {"ready": False} + + with patch('time.sleep'): # Speed up the test + with pytest.raises(RuntimeError, match="Gave up waiting for index creation"): + await processor.store_document_embeddings(message) + + @pytest.mark.asyncio + async def test_store_document_embeddings_unicode_content(self, processor): + """Test storing document embeddings with Unicode content""" + message = MagicMock() + message.metadata = MagicMock() + message.metadata.user = 'test_user' + message.metadata.collection = 'test_collection' + + chunk = ChunkEmbeddings( + chunk="Document with Unicode: éñ中文🚀".encode('utf-8'), + vectors=[[0.1, 0.2, 0.3]] + ) + message.chunks = [chunk] + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + processor.pinecone.has_index.return_value = True + + with patch('uuid.uuid4', return_value='test-id'): + await processor.store_document_embeddings(message) + + # Verify Unicode content was properly decoded and stored + call_args = mock_index.upsert.call_args + stored_doc = call_args[1]['vectors'][0]['metadata']['doc'] + assert stored_doc == "Document with Unicode: éñ中文🚀" + + @pytest.mark.asyncio + async def test_store_document_embeddings_large_chunks(self, processor): + """Test storing document embeddings with large document chunks""" + message = MagicMock() + message.metadata = MagicMock() + message.metadata.user = 'test_user' + message.metadata.collection = 'test_collection' + + # Create a large document chunk + large_content = "A" * 10000 # 10KB of content + chunk = ChunkEmbeddings( + chunk=large_content.encode('utf-8'), + vectors=[[0.1, 0.2, 0.3]] + ) + message.chunks = [chunk] + + mock_index = MagicMock() + processor.pinecone.Index.return_value = mock_index + processor.pinecone.has_index.return_value = True + + with patch('uuid.uuid4', return_value='test-id'): + await processor.store_document_embeddings(message) + + # Verify large content was stored + call_args = mock_index.upsert.call_args + stored_doc = call_args[1]['vectors'][0]['metadata']['doc'] + assert stored_doc == large_content + + def test_add_args_method(self): + """Test that add_args properly configures argument parser""" + from argparse import ArgumentParser + from unittest.mock import patch + + parser = ArgumentParser() + + # Mock the parent class add_args method + with patch('trustgraph.storage.doc_embeddings.pinecone.write.DocumentEmbeddingsStoreService.add_args') as mock_parent_add_args: + Processor.add_args(parser) + + # Verify parent add_args was called + mock_parent_add_args.assert_called_once() + + # Verify our specific arguments were added + args = parser.parse_args([]) + + assert hasattr(args, 'api_key') + assert args.api_key == 'not-specified' # Default value when no env var + assert hasattr(args, 'url') + assert args.url is None + assert hasattr(args, 'cloud') + assert args.cloud == 'aws' + assert hasattr(args, 'region') + assert args.region == 'us-east-1' + + def test_add_args_with_custom_values(self): + """Test add_args with custom command line values""" + from argparse import ArgumentParser + from unittest.mock import patch + + parser = ArgumentParser() + + with patch('trustgraph.storage.doc_embeddings.pinecone.write.DocumentEmbeddingsStoreService.add_args'): + Processor.add_args(parser) + + # Test parsing with custom values + args = parser.parse_args([ + '--api-key', 'custom-api-key', + '--url', 'https://custom-host.pinecone.io', + '--cloud', 'gcp', + '--region', 'us-west1' + ]) + + assert args.api_key == 'custom-api-key' + assert args.url == 'https://custom-host.pinecone.io' + assert args.cloud == 'gcp' + assert args.region == 'us-west1' + + def test_add_args_short_form(self): + """Test add_args with short form arguments""" + from argparse import ArgumentParser + from unittest.mock import patch + + parser = ArgumentParser() + + with patch('trustgraph.storage.doc_embeddings.pinecone.write.DocumentEmbeddingsStoreService.add_args'): + Processor.add_args(parser) + + # Test parsing with short form + args = parser.parse_args([ + '-a', 'short-api-key', + '-u', 'https://short-host.pinecone.io' + ]) + + assert args.api_key == 'short-api-key' + assert args.url == 'https://short-host.pinecone.io' + + @patch('trustgraph.storage.doc_embeddings.pinecone.write.Processor.launch') + def test_run_function(self, mock_launch): + """Test the run function calls Processor.launch with correct parameters""" + from trustgraph.storage.doc_embeddings.pinecone.write import run, default_ident + + run() + + mock_launch.assert_called_once_with( + default_ident, + "\nAccepts document chunks/vector pairs and writes them to a Pinecone store.\n" + ) \ No newline at end of file diff --git a/trustgraph-flow/trustgraph/query/doc_embeddings/pinecone/service.py b/trustgraph-flow/trustgraph/query/doc_embeddings/pinecone/service.py index 74c52055..8388a8ca 100755 --- a/trustgraph-flow/trustgraph/query/doc_embeddings/pinecone/service.py +++ b/trustgraph-flow/trustgraph/query/doc_embeddings/pinecone/service.py @@ -10,30 +10,21 @@ from pinecone.grpc import PineconeGRPC, GRPCClientConfig import uuid import os -from .... schema import DocumentEmbeddingsRequest, DocumentEmbeddingsResponse -from .... schema import Error, Value -from .... schema import document_embeddings_request_queue -from .... schema import document_embeddings_response_queue -from .... base import ConsumerProducer +from .... base import DocumentEmbeddingsQueryService -module = "de-query" - -default_input_queue = document_embeddings_request_queue -default_output_queue = document_embeddings_response_queue -default_subscriber = module +default_ident = "de-query" default_api_key = os.getenv("PINECONE_API_KEY", "not-specified") -class Processor(ConsumerProducer): +class Processor(DocumentEmbeddingsQueryService): 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) - self.url = params.get("url", None) self.api_key = params.get("api_key", default_api_key) + if self.api_key is None or self.api_key == "not-specified": + raise RuntimeError("Pinecone API key must be specified") + if self.url: self.pinecone = PineconeGRPC( @@ -47,88 +38,53 @@ class Processor(ConsumerProducer): super(Processor, self).__init__( **params | { - "input_queue": input_queue, - "output_queue": output_queue, - "subscriber": subscriber, - "input_schema": DocumentEmbeddingsRequest, - "output_schema": DocumentEmbeddingsResponse, "url": self.url, + "api_key": self.api_key, } ) - async def handle(self, msg): + async def query_document_embeddings(self, msg): try: - v = msg.value() - - # Sender-produced ID - id = msg.properties()["id"] - - print(f"Handling input {id}...", flush=True) + # Handle zero limit case + if msg.limit <= 0: + return [] chunks = [] - for vec in v.vectors: + for vec in msg.vectors: dim = len(vec) index_name = ( - "d-" + v.user + "-" + str(dim) + "d-" + msg.user + "-" + msg.collection + "-" + str(dim) ) index = self.pinecone.Index(index_name) results = index.query( - namespace=v.collection, vector=vec, - top_k=v.limit, + top_k=msg.limit, include_values=False, include_metadata=True ) - search_result = self.client.query_points( - collection_name=collection, - query=vec, - limit=v.limit, - with_payload=True, - ).points - for r in results.matches: doc = r.metadata["doc"] - chunks.add(doc) + chunks.append(doc) - print("Send response...", flush=True) - r = DocumentEmbeddingsResponse(documents=chunks, error=None) - await self.send(r, properties={"id": id}) - - print("Done.", flush=True) + return chunks except Exception as e: print(f"Exception: {e}") - - print("Send error response...", flush=True) - - r = DocumentEmbeddingsResponse( - error=Error( - type = "llm-error", - message = str(e), - ), - documents=None, - ) - - await self.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + raise e @staticmethod def add_args(parser): - ConsumerProducer.add_args( - parser, default_input_queue, default_subscriber, - default_output_queue, - ) + DocumentEmbeddingsQueryService.add_args(parser) parser.add_argument( '-a', '--api-key', @@ -143,5 +99,5 @@ class Processor(ConsumerProducer): def run(): - Processor.launch(module, __doc__) + Processor.launch(default_ident, __doc__) diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py index 128323aa..0d8bac83 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py @@ -1,42 +1,32 @@ """ -Accepts entity/vector pairs and writes them to a Qdrant store. +Accepts document chunks/vector pairs and writes them to a Pinecone store. """ -from qdrant_client import QdrantClient -from qdrant_client.models import PointStruct -from qdrant_client.models import Distance, VectorParams +from pinecone import Pinecone, ServerlessSpec +from pinecone.grpc import PineconeGRPC, GRPCClientConfig import time import uuid import os -from .... schema import DocumentEmbeddings -from .... schema import document_embeddings_store_queue -from .... log_level import LogLevel -from .... base import Consumer +from .... base import DocumentEmbeddingsStoreService -module = "de-write" - -default_input_queue = document_embeddings_store_queue -default_subscriber = module +default_ident = "de-write" default_api_key = os.getenv("PINECONE_API_KEY", "not-specified") default_cloud = "aws" default_region = "us-east-1" -class Processor(Consumer): +class Processor(DocumentEmbeddingsStoreService): def __init__(self, **params): - input_queue = params.get("input_queue", default_input_queue) - subscriber = params.get("subscriber", default_subscriber) - self.url = params.get("url", None) self.cloud = params.get("cloud", default_cloud) self.region = params.get("region", default_region) self.api_key = params.get("api_key", default_api_key) - if self.api_key is None: + if self.api_key is None or self.api_key == "not-specified": raise RuntimeError("Pinecone API key must be specified") if self.url: @@ -52,94 +42,96 @@ class Processor(Consumer): super(Processor, self).__init__( **params | { - "input_queue": input_queue, - "subscriber": subscriber, - "input_schema": DocumentEmbeddings, "url": self.url, + "cloud": self.cloud, + "region": self.region, + "api_key": self.api_key, } ) self.last_index_name = None - async def handle(self, msg): + def create_index(self, index_name, dim): - v = msg.value() + self.pinecone.create_index( + name = index_name, + dimension = dim, + metric = "cosine", + spec = ServerlessSpec( + cloud = self.cloud, + region = self.region, + ) + ) - for emb in v.chunks: + for i in range(0, 1000): + if self.pinecone.describe_index( + index_name + ).status["ready"]: + break + + time.sleep(1) + + if not self.pinecone.describe_index( + index_name + ).status["ready"]: + raise RuntimeError( + "Gave up waiting for index creation" + ) + + async def store_document_embeddings(self, message): + + for emb in message.chunks: + + if emb.chunk is None or emb.chunk == b"": continue + chunk = emb.chunk.decode("utf-8") - if chunk == "" or chunk is None: continue + if chunk == "": continue for vec in emb.vectors: - for vec in v.vectors: + dim = len(vec) + index_name = ( + "d-" + message.metadata.user + "-" + message.metadata.collection + "-" + str(dim) + ) - dim = len(vec) - collection = ( - "d-" + v.metadata.user + "-" + str(dim) - ) + if index_name != self.last_index_name: - if index_name != self.last_index_name: + if not self.pinecone.has_index(index_name): - if not self.pinecone.has_index(index_name): + try: - try: + self.create_index(index_name, dim) - self.pinecone.create_index( - name = index_name, - dimension = dim, - metric = "cosine", - spec = ServerlessSpec( - cloud = self.cloud, - region = self.region, - ) - ) + except Exception as e: + print("Pinecone index creation failed") + raise e - for i in range(0, 1000): + print(f"Index {index_name} created", flush=True) - if self.pinecone.describe_index( - index_name - ).status["ready"]: - break + self.last_index_name = index_name - time.sleep(1) + index = self.pinecone.Index(index_name) - if not self.pinecone.describe_index( - index_name - ).status["ready"]: - raise RuntimeError( - "Gave up waiting for index creation" - ) + # Generate unique ID for each vector + vector_id = str(uuid.uuid4()) - except Exception as e: - print("Pinecone index creation failed") - raise e + records = [ + { + "id": vector_id, + "values": vec, + "metadata": { "doc": chunk }, + } + ] - print(f"Index {index_name} created", flush=True) - - self.last_index_name = index_name - - index = self.pinecone.Index(index_name) - - records = [ - { - "id": id, - "values": vec, - "metadata": { "doc": chunk }, - } - ] - - index.upsert( - vectors = records, - namespace = v.metadata.collection, - ) + index.upsert( + vectors = records, + ) @staticmethod def add_args(parser): - Consumer.add_args( - parser, default_input_queue, default_subscriber, - ) + DocumentEmbeddingsStoreService.add_args(parser) parser.add_argument( '-a', '--api-key', @@ -166,5 +158,5 @@ class Processor(Consumer): def run(): - Processor.launch(module, __doc__) + Processor.launch(default_ident, __doc__)