From b0d84b7b9b590ebcffd4633283016ef6f87adce1 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 23 Feb 2026 19:56:00 +0000 Subject: [PATCH] Row embeddings agent tool --- trustgraph-base/trustgraph/base/__init__.py | 1 + .../base/row_embeddings_query_client.py | 45 +++++++++++++++ .../trustgraph/agent/react/service.py | 34 ++++++++++-- .../trustgraph/agent/react/tools.py | 55 +++++++++++++++++++ 4 files changed, 130 insertions(+), 5 deletions(-) create mode 100644 trustgraph-base/trustgraph/base/row_embeddings_query_client.py diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index e8530f6c..557109a2 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/trustgraph-base/trustgraph/base/__init__.py @@ -34,5 +34,6 @@ from . tool_service import ToolService from . tool_client import ToolClientSpec from . agent_client import AgentClientSpec from . structured_query_client import StructuredQueryClientSpec +from . row_embeddings_query_client import RowEmbeddingsQueryClientSpec from . collection_config_handler import CollectionConfigHandler diff --git a/trustgraph-base/trustgraph/base/row_embeddings_query_client.py b/trustgraph-base/trustgraph/base/row_embeddings_query_client.py new file mode 100644 index 00000000..0141da31 --- /dev/null +++ b/trustgraph-base/trustgraph/base/row_embeddings_query_client.py @@ -0,0 +1,45 @@ +from . request_response_spec import RequestResponse, RequestResponseSpec +from .. schema import RowEmbeddingsRequest, RowEmbeddingsResponse + +class RowEmbeddingsQueryClient(RequestResponse): + async def row_embeddings_query( + self, vectors, schema_name, user="trustgraph", collection="default", + index_name=None, limit=10, timeout=600 + ): + request = RowEmbeddingsRequest( + vectors=vectors, + schema_name=schema_name, + user=user, + collection=collection, + limit=limit + ) + if index_name: + request.index_name = index_name + + resp = await self.request(request, timeout=timeout) + + if resp.error: + raise RuntimeError(resp.error.message) + + # Return matches as list of dicts + return [ + { + "index_name": match.index_name, + "index_value": match.index_value, + "text": match.text, + "score": match.score + } + for match in (resp.matches or []) + ] + +class RowEmbeddingsQueryClientSpec(RequestResponseSpec): + def __init__( + self, request_name, response_name, + ): + super(RowEmbeddingsQueryClientSpec, self).__init__( + request_name = request_name, + request_schema = RowEmbeddingsRequest, + response_name = response_name, + response_schema = RowEmbeddingsResponse, + impl = RowEmbeddingsQueryClient, + ) diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py index 3af851d2..68bb5c14 100755 --- a/trustgraph-flow/trustgraph/agent/react/service.py +++ b/trustgraph-flow/trustgraph/agent/react/service.py @@ -13,10 +13,11 @@ logger = logging.getLogger(__name__) from ... base import AgentService, TextCompletionClientSpec, PromptClientSpec from ... base import GraphRagClientSpec, ToolClientSpec, StructuredQueryClientSpec +from ... base import RowEmbeddingsQueryClientSpec, EmbeddingsClientSpec from ... schema import AgentRequest, AgentResponse, AgentStep, Error -from . tools import KnowledgeQueryImpl, TextCompletionImpl, McpToolImpl, PromptImpl, StructuredQueryImpl +from . tools import KnowledgeQueryImpl, TextCompletionImpl, McpToolImpl, PromptImpl, StructuredQueryImpl, RowEmbeddingsQueryImpl from . agent_manager import AgentManager from ..tool_filter import validate_tool_config, filter_tools_by_group_and_state, get_next_state @@ -87,6 +88,20 @@ class Processor(AgentService): ) ) + self.register_specification( + EmbeddingsClientSpec( + request_name = "embeddings-request", + response_name = "embeddings-response", + ) + ) + + self.register_specification( + RowEmbeddingsQueryClientSpec( + request_name = "row-embeddings-query-request", + response_name = "row-embeddings-query-response", + ) + ) + async def on_tools_config(self, config, version): logger.info(f"Loading configuration version {version}") @@ -147,11 +162,20 @@ class Processor(AgentService): ) elif impl_id == "structured-query": impl = functools.partial( - StructuredQueryImpl, + StructuredQueryImpl, collection=data.get("collection"), user=None # User will be provided dynamically via context ) arguments = StructuredQueryImpl.get_arguments() + elif impl_id == "row-embeddings-query": + impl = functools.partial( + RowEmbeddingsQueryImpl, + schema_name=data.get("schema-name"), + collection=data.get("collection"), + user=None, # User will be provided dynamically via context + index_name=data.get("index-name") # Optional filter + ) + arguments = RowEmbeddingsQueryImpl.get_arguments() else: raise RuntimeError( f"Tool type {impl_id} not known" @@ -327,11 +351,11 @@ class Processor(AgentService): def __init__(self, flow, user): self._flow = flow self._user = user - + def __call__(self, service_name): client = self._flow(service_name) - # For structured query clients, store user context - if service_name == "structured-query-request": + # For query clients that need user context, store it + if service_name in ("structured-query-request", "row-embeddings-query-request"): client._current_user = self._user return client diff --git a/trustgraph-flow/trustgraph/agent/react/tools.py b/trustgraph-flow/trustgraph/agent/react/tools.py index e32dc2d8..c77929d9 100644 --- a/trustgraph-flow/trustgraph/agent/react/tools.py +++ b/trustgraph-flow/trustgraph/agent/react/tools.py @@ -128,6 +128,61 @@ class StructuredQueryImpl: return str(result) +# This tool implementation knows how to query row embeddings for semantic search +class RowEmbeddingsQueryImpl: + def __init__(self, context, schema_name, collection=None, user=None, index_name=None): + self.context = context + self.schema_name = schema_name + self.collection = collection + self.user = user + self.index_name = index_name # Optional: filter to specific index + + @staticmethod + def get_arguments(): + return [ + Argument( + name="query", + type="string", + description="Text to search for semantically similar values in the structured data index" + ) + ] + + async def invoke(self, **arguments): + # First get embeddings for the query text + embeddings_client = self.context("embeddings-request") + logger.debug("Getting embeddings for row query...") + + query_text = arguments.get("query") + vectors = await embeddings_client.embed(query_text) + + # Now query row embeddings + client = self.context("row-embeddings-query-request") + logger.debug("Row embeddings query...") + + # Get user from client context if available + user = getattr(client, '_current_user', self.user or "trustgraph") + + matches = await client.row_embeddings_query( + vectors=vectors, + schema_name=self.schema_name, + user=user, + collection=self.collection or "default", + index_name=self.index_name, + limit=10 + ) + + # Format results for agent consumption + if not matches: + return "No matching records found" + + results = [] + for match in matches: + result = f"- {match['index_name']}: {', '.join(match['index_value'])} (score: {match['score']:.3f})" + results.append(result) + + return "Matching records:\n" + "\n".join(results) + + # This tool implementation knows how to execute prompt templates class PromptImpl: def __init__(self, context, template_id, arguments=None):