From 6efdf185552a5a3433561454b005e30f13848d8a Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Thu, 4 Sep 2025 16:18:42 +0100
Subject: [PATCH] Agent integration to structured query

---
 trustgraph-base/trustgraph/base/__init__.py   |  1 +
 .../base/structured_query_client.py           | 33 +++++++++++++++++
 .../trustgraph/agent/react/service.py         | 17 ++++++++-
 .../trustgraph/agent/react/tools.py           | 37 +++++++++++++++++++
 4 files changed, 86 insertions(+), 2 deletions(-)
 create mode 100644 trustgraph-base/trustgraph/base/structured_query_client.py

diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py
index 5e279c8e..7ef199d3 100644
--- a/trustgraph-base/trustgraph/base/__init__.py
+++ b/trustgraph-base/trustgraph/base/__init__.py
@@ -31,4 +31,5 @@ from . graph_rag_client import GraphRagClientSpec
 from . tool_service import ToolService
 from . tool_client import ToolClientSpec
 from . agent_client import AgentClientSpec
+from . structured_query_client import StructuredQueryClientSpec
 
diff --git a/trustgraph-base/trustgraph/base/structured_query_client.py b/trustgraph-base/trustgraph/base/structured_query_client.py
new file mode 100644
index 00000000..dc025c4a
--- /dev/null
+++ b/trustgraph-base/trustgraph/base/structured_query_client.py
@@ -0,0 +1,33 @@
+from . request_response_spec import RequestResponse, RequestResponseSpec
+from .. schema import StructuredQueryRequest, StructuredQueryResponse
+
+class StructuredQueryClient(RequestResponse):
+    async def structured_query(self, question, timeout=600):
+        resp = await self.request(
+            StructuredQueryRequest(
+                question = question
+            ),
+            timeout=timeout
+        )
+
+        if resp.error:
+            raise RuntimeError(resp.error.message)
+
+        # Return the full response structure for the tool to handle
+        return {
+            "data": resp.data,
+            "errors": resp.errors if resp.errors else [],
+            "error": resp.error
+        }
+
+class StructuredQueryClientSpec(RequestResponseSpec):
+    def __init__(
+            self, request_name, response_name,
+    ):
+        super(StructuredQueryClientSpec, self).__init__(
+            request_name = request_name,
+            request_schema = StructuredQueryRequest,
+            response_name = response_name,
+            response_schema = StructuredQueryResponse,
+            impl = StructuredQueryClient,
+        )
\ No newline at end of file
diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py
index fbdb3f96..55f8ce45 100755
--- a/trustgraph-flow/trustgraph/agent/react/service.py
+++ b/trustgraph-flow/trustgraph/agent/react/service.py
@@ -12,11 +12,11 @@ import logging
 logger = logging.getLogger(__name__)
 
 from ... base import AgentService, TextCompletionClientSpec, PromptClientSpec
-from ... base import GraphRagClientSpec, ToolClientSpec
+from ... base import GraphRagClientSpec, ToolClientSpec, StructuredQueryClientSpec
 
 from ... schema import AgentRequest, AgentResponse, AgentStep, Error
 
-from . tools import KnowledgeQueryImpl, TextCompletionImpl, McpToolImpl, PromptImpl
+from . tools import KnowledgeQueryImpl, TextCompletionImpl, McpToolImpl, PromptImpl, StructuredQueryImpl
 from . agent_manager import AgentManager
 from ..tool_filter import validate_tool_config, filter_tools_by_group_and_state, get_next_state
 
@@ -80,6 +80,13 @@ class Processor(AgentService):
             )
         )
 
+        self.register_specification(
+            StructuredQueryClientSpec(
+                request_name = "structured-query-request",
+                response_name = "structured-query-response",
+            )
+        )
+
     async def on_tools_config(self, config, version):
 
         logger.info(f"Loading configuration version {version}")
@@ -138,6 +145,12 @@ class Processor(AgentService):
                             template_id=data.get("template"),
                             arguments=arguments
                         )
+                    elif impl_id == "structured-query":
+                        impl = functools.partial(
+                            StructuredQueryImpl, 
+                            collection=data.get("collection")
+                        )
+                        arguments = StructuredQueryImpl.get_arguments()
                     else:
                         raise RuntimeError(
                             f"Tool type {impl_id} not known"
diff --git a/trustgraph-flow/trustgraph/agent/react/tools.py b/trustgraph-flow/trustgraph/agent/react/tools.py
index 948424ec..596df741 100644
--- a/trustgraph-flow/trustgraph/agent/react/tools.py
+++ b/trustgraph-flow/trustgraph/agent/react/tools.py
@@ -85,6 +85,43 @@ class McpToolImpl:
             return json.dumps(output)
 
 
+# This tool implementation knows how to query structured data using natural language
+class StructuredQueryImpl:
+    def __init__(self, context, collection=None):
+        self.context = context
+        self.collection = collection  # For multi-tenant scenarios
+    
+    @staticmethod
+    def get_arguments():
+        return [
+            Argument(
+                name="question",
+                type="string", 
+                description="Natural language question about structured data (tables, databases, etc.)"
+            )
+        ]
+    
+    async def invoke(self, **arguments):
+        client = self.context("structured-query-request")
+        logger.debug("Structured query question...")
+        
+        result = await client.structured_query(
+            arguments.get("question")
+        )
+        
+        # Format the result for the agent
+        if isinstance(result, dict):
+            if result.get("error"):
+                return f"Error: {result['error']['message']}"
+            elif result.get("data"):
+                # Pretty format JSON data for agent consumption
+                return json.dumps(result["data"], indent=2)
+            else:
+                return "No data returned"
+        else:
+            return str(result)
+
+
 # This tool implementation knows how to execute prompt templates
 class PromptImpl:
     def __init__(self, context, template_id, arguments=None):