Row embeddings agent tool

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
 

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,
+        )

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
 

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):