diff --git a/tests/unit/test_cores/test_knowledge_manager.py b/tests/unit/test_cores/test_knowledge_manager.py index 96c9c427..8c37ca32 100644 --- a/tests/unit/test_cores/test_knowledge_manager.py +++ b/tests/unit/test_cores/test_knowledge_manager.py @@ -98,7 +98,7 @@ def sample_graph_embeddings(): entities=[ EntityEmbeddings( entity=Term(type=IRI, iri="http://example.org/john"), - vectors=[[0.1, 0.2, 0.3]] + vector=[0.1, 0.2, 0.3] ) ] ) diff --git a/tests/unit/test_storage/test_doc_embeddings_pinecone_storage.py b/tests/unit/test_storage/test_doc_embeddings_pinecone_storage.py index e5c1e91e..fec4f87e 100644 --- a/tests/unit/test_storage/test_doc_embeddings_pinecone_storage.py +++ b/tests/unit/test_storage/test_doc_embeddings_pinecone_storage.py @@ -286,37 +286,39 @@ class TestPineconeDocEmbeddingsStorageProcessor: 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 - ] + # Each chunk has a single vector of different dimensions + chunk1 = ChunkEmbeddings( + chunk=b"Document chunk 1", + vector=[0.1, 0.2] # 2D vector ) - message.chunks = [chunk] - - mock_index_2d = MagicMock() - mock_index_4d = MagicMock() - mock_index_3d = MagicMock() - + chunk2 = ChunkEmbeddings( + chunk=b"Document chunk 2", + vector=[0.3, 0.4, 0.5, 0.6] # 4D vector + ) + chunk3 = ChunkEmbeddings( + chunk=b"Document chunk 3", + vector=[0.7, 0.8, 0.9] # 3D vector + ) + message.chunks = [chunk1, chunk2, chunk3] + + mock_index = MagicMock() + def mock_index_side_effect(name): # All dimensions now use the same index name pattern - # Different dimensions will be handled within the same index if "test_user" in name and "test_collection" in name: - return mock_index_2d # Just return one mock for all + return mock_index return MagicMock() - + 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 all vectors are now stored in the same index - # (Pinecone can handle mixed dimensions in the same index) - assert processor.pinecone.Index.call_count == 3 # Called once per vector - mock_index_2d.upsert.call_count == 3 # All upserts go to same index + # (Each chunk has a single vector, called once per chunk) + assert processor.pinecone.Index.call_count == 3 # Called once per chunk + assert mock_index.upsert.call_count == 3 # All upserts go to same index @pytest.mark.asyncio async def test_store_document_embeddings_empty_chunks_list(self, processor): @@ -346,7 +348,7 @@ class TestPineconeDocEmbeddingsStorageProcessor: chunk = ChunkEmbeddings( chunk=b"Document with no vectors", - vectors=[] + vector=[] ) message.chunks = [chunk] diff --git a/tests/unit/test_storage/test_graph_embeddings_pinecone_storage.py b/tests/unit/test_storage/test_graph_embeddings_pinecone_storage.py index 0fd0fde3..bba30caa 100644 --- a/tests/unit/test_storage/test_graph_embeddings_pinecone_storage.py +++ b/tests/unit/test_storage/test_graph_embeddings_pinecone_storage.py @@ -24,16 +24,20 @@ class TestPineconeGraphEmbeddingsStorageProcessor: message.metadata.user = 'test_user' message.metadata.collection = 'test_collection' - # Create test entity embeddings + # Create test entity embeddings (each entity has a single vector) entity1 = EntityEmbeddings( entity=Value(value="http://example.org/entity1", is_uri=True), - vectors=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + vector=[0.1, 0.2, 0.3] ) entity2 = EntityEmbeddings( - entity=Value(value="entity2", is_uri=False), - vectors=[[0.7, 0.8, 0.9]] + entity=Value(value="http://example.org/entity2", is_uri=True), + vector=[0.4, 0.5, 0.6] ) - message.entities = [entity1, entity2] + entity3 = EntityEmbeddings( + entity=Value(value="entity3", is_uri=False), + vector=[0.7, 0.8, 0.9] + ) + message.entities = [entity1, entity2, entity3] return message @@ -190,7 +194,7 @@ class TestPineconeGraphEmbeddingsStorageProcessor: entity = EntityEmbeddings( entity=Value(value="test_entity", is_uri=False), - vectors=[[0.1, 0.2, 0.3]] + vector=[0.1, 0.2, 0.3] ) message.entities = [entity] @@ -222,7 +226,7 @@ class TestPineconeGraphEmbeddingsStorageProcessor: entity = EntityEmbeddings( entity=Value(value="", is_uri=False), - vectors=[[0.1, 0.2, 0.3]] + vector=[0.1, 0.2, 0.3] ) message.entities = [entity] @@ -244,7 +248,7 @@ class TestPineconeGraphEmbeddingsStorageProcessor: entity = EntityEmbeddings( entity=Value(value=None, is_uri=False), - vectors=[[0.1, 0.2, 0.3]] + vector=[0.1, 0.2, 0.3] ) message.entities = [entity] @@ -322,7 +326,7 @@ class TestPineconeGraphEmbeddingsStorageProcessor: entity = EntityEmbeddings( entity=Value(value="test_entity", is_uri=False), - vectors=[] + vector=[] ) message.entities = [entity] @@ -344,7 +348,7 @@ class TestPineconeGraphEmbeddingsStorageProcessor: entity = EntityEmbeddings( entity=Value(value="test_entity", is_uri=False), - vectors=[[0.1, 0.2, 0.3]] + vector=[0.1, 0.2, 0.3] ) message.entities = [entity] @@ -369,7 +373,7 @@ class TestPineconeGraphEmbeddingsStorageProcessor: entity = EntityEmbeddings( entity=Value(value="test_entity", is_uri=False), - vectors=[[0.1, 0.2, 0.3]] + vector=[0.1, 0.2, 0.3] ) message.entities = [entity]