From 3d4f859aa2b9fc9d1e2088141f68e9b8e8ad9506 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Sun, 8 Mar 2026 18:34:22 +0000 Subject: [PATCH] Implementation --- .../test_row_embeddings_processor.py | 18 ++++++++++--- trustgraph-base/trustgraph/api/flow.py | 17 ++++++------ .../trustgraph/api/socket_client.py | 12 ++++----- .../messaging/translators/embeddings.py | 8 +++--- .../trustgraph/cli/invoke_embeddings.py | 19 +++++++++----- .../document_embeddings/embeddings.py | 6 +++-- .../embeddings/graph_embeddings/embeddings.py | 26 ++++++++++--------- .../embeddings/row_embeddings/embeddings.py | 14 +++++++--- 8 files changed, 73 insertions(+), 47 deletions(-) diff --git a/tests/unit/test_embeddings/test_row_embeddings_processor.py b/tests/unit/test_embeddings/test_row_embeddings_processor.py index 47405431..45a22e48 100644 --- a/tests/unit/test_embeddings/test_row_embeddings_processor.py +++ b/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() diff --git a/trustgraph-base/trustgraph/api/flow.py b/trustgraph-base/trustgraph/api/flow.py index c50bf9c4..5142aac4 100644 --- a/trustgraph-base/trustgraph/api/flow.py +++ b/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( diff --git a/trustgraph-base/trustgraph/api/socket_client.py b/trustgraph-base/trustgraph/api/socket_client.py index b471b535..e5d5a356 100644 --- a/trustgraph-base/trustgraph/api/socket_client.py +++ b/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) diff --git a/trustgraph-base/trustgraph/messaging/translators/embeddings.py b/trustgraph-base/trustgraph/messaging/translators/embeddings.py index 7e6eff83..454ce733 100644 --- a/trustgraph-base/trustgraph/messaging/translators/embeddings.py +++ b/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 } diff --git a/trustgraph-cli/trustgraph/cli/invoke_embeddings.py b/trustgraph-cli/trustgraph/cli/invoke_embeddings.py index 71a88bd7..699a85cf 100644 --- a/trustgraph-cli/trustgraph/cli/invoke_embeddings.py +++ b/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, ) diff --git a/trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py b/trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py index 032e15c4..eb21d418 100755 --- a/trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py +++ b/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( diff --git a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py index fec528bd..c54e719d 100755 --- a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py +++ b/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): diff --git a/trustgraph-flow/trustgraph/embeddings/row_embeddings/embeddings.py b/trustgraph-flow/trustgraph/embeddings/row_embeddings/embeddings.py index 84c41ff3..c1d04302 100644 --- a/trustgraph-flow/trustgraph/embeddings/row_embeddings/embeddings.py +++ b/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 ) )