Fixing tests

tests/unit/test_storage/test_doc_embeddings_milvus_storage.py

@@ -245,26 +245,31 @@ class TestMilvusDocEmbeddingsStorageProcessor:
         message.metadata = MagicMock()
         message.metadata.user = 'test_user'
         message.metadata.collection = 'test_collection'
-        
+
-        chunk = ChunkEmbeddings(
+        # Each chunk has a single vector of different dimensions
-            chunk_id="Document with mixed dimensions",
+        chunk1 = ChunkEmbeddings(
-            vectors=[
+            chunk_id="chunk/doc/2d",
-                [0.1, 0.2],  # 2D vector
+            vector=[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]
+        chunk2 = ChunkEmbeddings(
-        
+            chunk_id="chunk/doc/4d",
+            vector=[0.3, 0.4, 0.5, 0.6]  # 4D vector
+        )
+        chunk3 = ChunkEmbeddings(
+            chunk_id="chunk/doc/3d",
+            vector=[0.7, 0.8, 0.9]  # 3D vector
+        )
+        message.chunks = [chunk1, chunk2, chunk3]
+
         await processor.store_document_embeddings(message)
-        
+
         # Verify all vectors were inserted regardless of dimension with user/collection parameters
         expected_calls = [
-            ([0.1, 0.2], "Document with mixed dimensions", 'test_user', 'test_collection'),
+            ([0.1, 0.2], "chunk/doc/2d", 'test_user', 'test_collection'),
-            ([0.3, 0.4, 0.5, 0.6], "Document with mixed dimensions", 'test_user', 'test_collection'),
+            ([0.3, 0.4, 0.5, 0.6], "chunk/doc/4d", 'test_user', 'test_collection'),
-            ([0.7, 0.8, 0.9], "Document with mixed dimensions", 'test_user', 'test_collection'),
+            ([0.7, 0.8, 0.9], "chunk/doc/3d", 'test_user', 'test_collection'),
         ]
-        
+
         assert processor.vecstore.insert.call_count == 3
         for i, (expected_vec, expected_doc, expected_user, expected_collection) in enumerate(expected_calls):
             actual_call = processor.vecstore.insert.call_args_list[i]

tests/unit/test_storage/test_doc_embeddings_pinecone_storage.py

@@ -27,11 +27,11 @@ class TestPineconeDocEmbeddingsStorageProcessor:
         # 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]]
+            vector=[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]]
+            vector=[0.7, 0.8, 0.9]
         )
         message.chunks = [chunk1, chunk2]
         
@@ -125,7 +125,7 @@ class TestPineconeDocEmbeddingsStorageProcessor:
         
         chunk = ChunkEmbeddings(
             chunk=b"Test document content",
-            vectors=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
+            vector=[0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
         )
         message.chunks = [chunk]
         
@@ -190,7 +190,7 @@ class TestPineconeDocEmbeddingsStorageProcessor:
 
         chunk = ChunkEmbeddings(
             chunk=b"Test document content",
-            vectors=[[0.1, 0.2, 0.3]]
+            vector=[0.1, 0.2, 0.3]
         )
         message.chunks = [chunk]
 
@@ -222,7 +222,7 @@ class TestPineconeDocEmbeddingsStorageProcessor:
         
         chunk = ChunkEmbeddings(
             chunk=b"",
-            vectors=[[0.1, 0.2, 0.3]]
+            vector=[0.1, 0.2, 0.3]
         )
         message.chunks = [chunk]
         
@@ -244,7 +244,7 @@ class TestPineconeDocEmbeddingsStorageProcessor:
         
         chunk = ChunkEmbeddings(
             chunk=None,
-            vectors=[[0.1, 0.2, 0.3]]
+            vector=[0.1, 0.2, 0.3]
         )
         message.chunks = [chunk]
         
@@ -266,7 +266,7 @@ class TestPineconeDocEmbeddingsStorageProcessor:
         
         chunk = ChunkEmbeddings(
             chunk=b"",  # Empty bytes
-            vectors=[[0.1, 0.2, 0.3]]
+            vector=[0.1, 0.2, 0.3]
         )
         message.chunks = [chunk]
         
@@ -368,7 +368,7 @@ class TestPineconeDocEmbeddingsStorageProcessor:
 
         chunk = ChunkEmbeddings(
             chunk=b"Test document content",
-            vectors=[[0.1, 0.2, 0.3]]
+            vector=[0.1, 0.2, 0.3]
         )
         message.chunks = [chunk]
 
@@ -393,7 +393,7 @@ class TestPineconeDocEmbeddingsStorageProcessor:
 
         chunk = ChunkEmbeddings(
             chunk=b"Test document content",
-            vectors=[[0.1, 0.2, 0.3]]
+            vector=[0.1, 0.2, 0.3]
         )
         message.chunks = [chunk]
 
@@ -419,7 +419,7 @@ class TestPineconeDocEmbeddingsStorageProcessor:
         
         chunk = ChunkEmbeddings(
             chunk="Document with Unicode: éñ中文🚀".encode('utf-8'),
-            vectors=[[0.1, 0.2, 0.3]]
+            vector=[0.1, 0.2, 0.3]
         )
         message.chunks = [chunk]
         
@@ -447,7 +447,7 @@ class TestPineconeDocEmbeddingsStorageProcessor:
         large_content = "A" * 10000  # 10KB of content
         chunk = ChunkEmbeddings(
             chunk=large_content.encode('utf-8'),
-            vectors=[[0.1, 0.2, 0.3]]
+            vector=[0.1, 0.2, 0.3]
         )
         message.chunks = [chunk]
         

tests/unit/test_storage/test_graph_embeddings_milvus_storage.py

@@ -180,7 +180,7 @@ class TestMilvusGraphEmbeddingsStorageProcessor:
         )
         empty_entity = EntityEmbeddings(
             entity=Term(type=LITERAL, value=''),
-            vectors=[[0.4, 0.5, 0.6]],
+            vector=[0.4, 0.5, 0.6],
             chunk_id=''
         )
         none_entity = EntityEmbeddings(
@@ -238,26 +238,31 @@ class TestMilvusGraphEmbeddingsStorageProcessor:
         message.metadata = MagicMock()
         message.metadata.user = 'test_user'
         message.metadata.collection = 'test_collection'
-        
+
-        entity = EntityEmbeddings(
+        # Each entity has a single vector of different dimensions
-            entity=Term(type=IRI, iri='http://example.com/entity'),
+        entity1 = EntityEmbeddings(
-            vectors=[
+            entity=Term(type=IRI, iri='http://example.com/entity1'),
-                [0.1, 0.2],  # 2D vector
+            vector=[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.entities = [entity]
+        entity2 = EntityEmbeddings(
-        
+            entity=Term(type=IRI, iri='http://example.com/entity2'),
+            vector=[0.3, 0.4, 0.5, 0.6]  # 4D vector
+        )
+        entity3 = EntityEmbeddings(
+            entity=Term(type=IRI, iri='http://example.com/entity3'),
+            vector=[0.7, 0.8, 0.9]  # 3D vector
+        )
+        message.entities = [entity1, entity2, entity3]
+
         await processor.store_graph_embeddings(message)
-        
+
         # Verify all vectors were inserted regardless of dimension
         expected_calls = [
-            ([0.1, 0.2], 'http://example.com/entity'),
+            ([0.1, 0.2], 'http://example.com/entity1'),
-            ([0.3, 0.4, 0.5, 0.6], 'http://example.com/entity'),
+            ([0.3, 0.4, 0.5, 0.6], 'http://example.com/entity2'),
-            ([0.7, 0.8, 0.9], 'http://example.com/entity'),
+            ([0.7, 0.8, 0.9], 'http://example.com/entity3'),
         ]
-        
+
         assert processor.vecstore.insert.call_count == 3
         for i, (expected_vec, expected_entity) in enumerate(expected_calls):
             actual_call = processor.vecstore.insert.call_args_list[i]
@@ -278,7 +283,7 @@ class TestMilvusGraphEmbeddingsStorageProcessor:
         )
         literal_entity = EntityEmbeddings(
             entity=Term(type=LITERAL, value='literal entity text'),
-            vectors=[[0.4, 0.5, 0.6]]
+            vector=[0.4, 0.5, 0.6]
         )
         message.entities = [uri_entity, literal_entity]
         

tests/unit/test_storage/test_row_embeddings_qdrant_storage.py

@@ -197,7 +197,7 @@ class TestQdrantRowEmbeddingsStorage(IsolatedAsyncioTestCase):
             index_name='customer_id',
             index_value=['CUST001'],
             text='CUST001',
-            vectors=[[0.1, 0.2, 0.3]]
+            vector=[0.1, 0.2, 0.3]
         )
 
         embeddings_msg = RowEmbeddings(
@@ -255,7 +255,7 @@ class TestQdrantRowEmbeddingsStorage(IsolatedAsyncioTestCase):
             index_name='name',
             index_value=['John Doe'],
             text='John Doe',
-            vectors=[[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]
+            vector=[0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
         )
 
         embeddings_msg = RowEmbeddings(
@@ -299,7 +299,7 @@ class TestQdrantRowEmbeddingsStorage(IsolatedAsyncioTestCase):
             index_name='id',
             index_value=['123'],
             text='123',
-            vectors=[]  # Empty vectors
+            vector=[]  # Empty vector
         )
 
         embeddings_msg = RowEmbeddings(
@@ -342,7 +342,7 @@ class TestQdrantRowEmbeddingsStorage(IsolatedAsyncioTestCase):
             index_name='id',
             index_value=['123'],
             text='123',
-            vectors=[[0.1, 0.2]]
+            vector=[0.1, 0.2]
         )
 
         embeddings_msg = RowEmbeddings(