Implementation

tests/unit/test_embeddings/test_row_embeddings_processor.py

@@ -353,7 +353,14 @@ class TestRowEmbeddingsProcessor(IsolatedAsyncioTestCase):
 
         # Mock the flow
         mock_embeddings_request = AsyncMock()
-        mock_embeddings_request.embed.return_value = [[0.1, 0.2, 0.3]]
+        # Return batch of vector sets (one per text)
+        # 4 unique texts: CUST001, John Doe, CUST002, Jane Smith
+        mock_embeddings_request.embed.return_value = [
+            [[0.1, 0.2, 0.3]],  # vectors for text 1
+            [[0.2, 0.3, 0.4]],  # vectors for text 2
+            [[0.3, 0.4, 0.5]],  # vectors for text 3
+            [[0.4, 0.5, 0.6]],  # vectors for text 4
+        ]
 
         mock_output = AsyncMock()
 
@@ -368,9 +375,12 @@ class TestRowEmbeddingsProcessor(IsolatedAsyncioTestCase):
 
         await processor.on_message(mock_msg, MagicMock(), mock_flow)
 
-        # Should have called embed for each unique text
-        # 4 values: CUST001, John Doe, CUST002, Jane Smith
-        assert mock_embeddings_request.embed.call_count == 4
+        # Should have called embed once with all texts in a batch
+        assert mock_embeddings_request.embed.call_count == 1
+        # Verify it was called with a list of texts
+        call_args = mock_embeddings_request.embed.call_args
+        assert 'texts' in call_args.kwargs
+        assert len(call_args.kwargs['texts']) == 4
 
         # Should have sent output
         mock_output.send.assert_called()

trustgraph-base/trustgraph/api/flow.py

@@ -544,30 +544,29 @@ class FlowInstance:
             input
         )["response"]
 
-    def embeddings(self, text):
+    def embeddings(self, texts):
         """
-        Generate vector embeddings for text.
+        Generate vector embeddings for one or more texts.
 
-        Converts text into dense vector representations suitable for semantic
+        Converts texts into dense vector representations suitable for semantic
         search and similarity comparison.
 
         Args:
-            text: Input text to embed
+            texts: List of input texts to embed
 
         Returns:
-            list[float]: Vector embedding
+            list[list[list[float]]]: Vector embeddings, one set per input text
 
         Example:
             ```python
             flow = api.flow().id("default")
-            vectors = flow.embeddings("quantum computing")
-            print(f"Embedding dimension: {len(vectors)}")
+            vectors = flow.embeddings(["quantum computing"])
+            print(f"Embedding dimension: {len(vectors[0][0])}")
             ```
         """
 
-        # The input consists of a text block
         input = {
-            "text": text
+            "texts": texts
         }
 
         return self.request(

trustgraph-base/trustgraph/api/socket_client.py

@@ -712,27 +712,27 @@ class SocketFlowInstance:
 
         return self.client._send_request_sync("document-embeddings", self.flow_id, request, False)
 
-    def embeddings(self, text: str, **kwargs: Any) -> Dict[str, Any]:
+    def embeddings(self, texts: list, **kwargs: Any) -> Dict[str, Any]:
         """
-        Generate vector embeddings for text.
+        Generate vector embeddings for one or more texts.
 
         Args:
-            text: Input text to embed
+            texts: List of input texts to embed
             **kwargs: Additional parameters passed to the service
 
         Returns:
-            dict: Response containing vectors
+            dict: Response containing vectors (one set per input text)
 
         Example:
             ```python
             socket = api.socket()
             flow = socket.flow("default")
 
-            result = flow.embeddings("quantum computing")
+            result = flow.embeddings(["quantum computing"])
             vectors = result.get("vectors", [])
             ```
         """
-        request = {"text": text}
+        request = {"texts": texts}
         request.update(kwargs)
 
         return self.client._send_request_sync("embeddings", self.flow_id, request, False)

trustgraph-base/trustgraph/messaging/translators/embeddings.py

@@ -5,15 +5,15 @@ from .base import MessageTranslator
 
 class EmbeddingsRequestTranslator(MessageTranslator):
     """Translator for EmbeddingsRequest schema objects"""
-    
+
     def to_pulsar(self, data: Dict[str, Any]) -> EmbeddingsRequest:
         return EmbeddingsRequest(
-            text=data["text"]
+            texts=data["texts"]
         )
-    
+
     def from_pulsar(self, obj: EmbeddingsRequest) -> Dict[str, Any]:
         return {
-            "text": obj.text
+            "texts": obj.texts
         }
 
 

trustgraph-cli/trustgraph/cli/invoke_embeddings.py

@@ -10,7 +10,7 @@ from trustgraph.api import Api
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
 default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def query(url, flow_id, text, token=None):
+def query(url, flow_id, texts, token=None):
 
     # Create API client
     api = Api(url=url, token=token)
@@ -19,9 +19,14 @@ def query(url, flow_id, text, token=None):
 
     try:
         # Call embeddings service
-        result = flow.embeddings(text=text)
+        result = flow.embeddings(texts=texts)
         vectors = result.get("vectors", [])
-        print(vectors)
+        # Print each text's vectors
+        for i, vecs in enumerate(vectors):
+            if len(texts) > 1:
+                print(f"Text {i + 1}: {vecs}")
+            else:
+                print(vecs)
 
     finally:
         # Clean up socket connection
@@ -53,9 +58,9 @@ def main():
     )
 
     parser.add_argument(
-        'text',
-        nargs=1,
-        help='Text to convert to embedding vector',
+        'texts',
+        nargs='+',
+        help='Text(s) to convert to embedding vectors',
     )
 
     args = parser.parse_args()
@@ -65,7 +70,7 @@ def main():
         query(
             url=args.url,
             flow_id=args.flow_id,
-            text=args.text[0],
+            texts=args.texts,
             token=args.token,
         )
 

trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py

@@ -62,11 +62,13 @@ class Processor(FlowProcessor):
 
             resp = await flow("embeddings-request").request(
                 EmbeddingsRequest(
-                    text = v.chunk
+                    texts=[v.chunk]
                 )
             )
 
-            vectors = resp.vectors
+            # vectors[0] is the vector set for the first (only) text
+            # vectors[0][0] is the first vector in that set
+            vectors = resp.vectors[0][0] if resp.vectors else []
 
             embeds = [
                 ChunkEmbeddings(

trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py

@@ -58,23 +58,25 @@ class Processor(FlowProcessor):
         v = msg.value()
         logger.info(f"Indexing {v.metadata.id}...")
 
-        entities = []
-
         try:
 
-            for entity in v.entities:
+            # Collect all contexts for batch embedding
+            contexts = [entity.context for entity in v.entities]
 
-                vectors = await flow("embeddings-request").embed(
-                    text = entity.context
-                )
+            # Single batch embedding call
+            all_vectors = await flow("embeddings-request").embed(
+                texts=contexts
+            )
 
-                entities.append(
-                    EntityEmbeddings(
-                        entity=entity.entity,
-                        vectors=vectors,
-                        chunk_id=entity.chunk_id,  # Provenance: source chunk
-                    )
+            # Pair results with entities
+            entities = [
+                EntityEmbeddings(
+                    entity=entity.entity,
+                    vectors=vectors[0],  # First vector from the set
+                    chunk_id=entity.chunk_id,  # Provenance: source chunk
                 )
+                for entity, vectors in zip(v.entities, all_vectors)
+            ]
 
             # Send in batches to avoid oversized messages
             for i in range(0, len(entities), self.batch_size):

trustgraph-flow/trustgraph/embeddings/row_embeddings/embeddings.py

@@ -200,15 +200,23 @@ class Processor(CollectionConfigHandler, FlowProcessor):
         embeddings_list = []
 
         try:
-            for text, (index_name, index_value) in texts_to_embed.items():
-                vectors = await flow("embeddings-request").embed(text=text)
+            # Collect texts and metadata for batch embedding
+            texts = list(texts_to_embed.keys())
+            metadata = list(texts_to_embed.values())
 
+            # Single batch embedding call
+            all_vectors = await flow("embeddings-request").embed(texts=texts)
+
+            # Pair results with metadata
+            for text, (index_name, index_value), vectors in zip(
+                texts, metadata, all_vectors
+            ):
                 embeddings_list.append(
                     RowIndexEmbedding(
                         index_name=index_name,
                         index_value=index_value,
                         text=text,
-                        vectors=vectors
+                        vectors=vectors[0]  # First vector from the set
                     )
                 )