Fix/embeddings integration 2 (#670)

tests/integration/test_document_rag_integration.py

@@ -27,9 +27,13 @@ class TestDocumentRagIntegration:
     def mock_embeddings_client(self):
         """Mock embeddings client that returns realistic vector embeddings"""
         client = AsyncMock()
+        # New batch format: [[[vectors_for_text1], ...]]
+        # One text input returns one vector set containing two vectors
         client.embed.return_value = [
-            [0.1, 0.2, 0.3, 0.4, 0.5],  # Realistic 5-dimensional embedding
-            [0.6, 0.7, 0.8, 0.9, 1.0]   # Second embedding for testing
+            [
+                [0.1, 0.2, 0.3, 0.4, 0.5],  # First vector for text
+                [0.6, 0.7, 0.8, 0.9, 1.0]   # Second vector for text
+            ]
         ]
         return client
 
@@ -90,7 +94,7 @@ class TestDocumentRagIntegration:
         )
 
         # Assert - Verify service coordination
-        mock_embeddings_client.embed.assert_called_once_with(query)
+        mock_embeddings_client.embed.assert_called_once_with([query])
 
         mock_doc_embeddings_client.query.assert_called_once_with(
             [[0.1, 0.2, 0.3, 0.4, 0.5], [0.6, 0.7, 0.8, 0.9, 1.0]],

tests/integration/test_document_rag_streaming_integration.py

@@ -30,7 +30,8 @@ class TestDocumentRagStreaming:
     def mock_embeddings_client(self):
         """Mock embeddings client"""
         client = AsyncMock()
-        client.embed.return_value = [[0.1, 0.2, 0.3, 0.4, 0.5]]
+        # New batch format: [[[vectors_for_text1]]]
+        client.embed.return_value = [[[0.1, 0.2, 0.3, 0.4, 0.5]]]
         return client
 
     @pytest.fixture

tests/integration/test_graph_rag_integration.py

@@ -21,8 +21,12 @@ class TestGraphRagIntegration:
     def mock_embeddings_client(self):
         """Mock embeddings client that returns realistic vector embeddings"""
         client = AsyncMock()
+        # New batch format: [[[vectors_for_text1], ...]]
+        # One text input returns one vector set containing one vector
         client.embed.return_value = [
-            [0.1, 0.2, 0.3, 0.4, 0.5],  # Realistic 5-dimensional embedding
+            [
+                [0.1, 0.2, 0.3, 0.4, 0.5],  # Vector for text
+            ]
         ]
         return client
 
@@ -120,8 +124,8 @@ class TestGraphRagIntegration:
 
         # Assert - Verify service coordination
 
-        # 1. Should compute embeddings for query
-        mock_embeddings_client.embed.assert_called_once_with(query)
+        # 1. Should compute embeddings for query (now expects list of texts)
+        mock_embeddings_client.embed.assert_called_once_with([query])
 
         # 2. Should query graph embeddings to find relevant entities
         mock_graph_embeddings_client.query.assert_called_once()

tests/integration/test_graph_rag_streaming_integration.py

@@ -24,7 +24,8 @@ class TestGraphRagStreaming:
     def mock_embeddings_client(self):
         """Mock embeddings client"""
         client = AsyncMock()
-        client.embed.return_value = [[0.1, 0.2, 0.3, 0.4, 0.5]]
+        # New batch format: [[[vectors_for_text1]]]
+        client.embed.return_value = [[[0.1, 0.2, 0.3, 0.4, 0.5]]]
         return client
 
     @pytest.fixture

tests/integration/test_rag_streaming_protocol.py

@@ -18,7 +18,7 @@ class TestGraphRagStreamingProtocol:
     def mock_embeddings_client(self):
         """Mock embeddings client"""
         client = AsyncMock()
-        client.embed.return_value = [[0.1, 0.2, 0.3]]
+        client.embed.return_value = [[[0.1, 0.2, 0.3]]]
         return client
 
     @pytest.fixture
@@ -197,7 +197,7 @@ class TestDocumentRagStreamingProtocol:
     def mock_embeddings_client(self):
         """Mock embeddings client"""
         client = AsyncMock()
-        client.embed.return_value = [[0.1, 0.2, 0.3]]
+        client.embed.return_value = [[[0.1, 0.2, 0.3]]]
         return client
 
     @pytest.fixture