Tool service implementation

trustgraph-base/trustgraph/base/__init__.py

@@ -32,6 +32,7 @@ from . agent_service import AgentService
 from . graph_rag_client import GraphRagClientSpec
 from . tool_service import ToolService
 from . tool_client import ToolClientSpec
+from . tool_service_client import ToolServiceClientSpec
 from . agent_client import AgentClientSpec
 from . structured_query_client import StructuredQueryClientSpec
 from . row_embeddings_query_client import RowEmbeddingsQueryClientSpec

trustgraph-base/trustgraph/base/tool_service_client.py

@@ -0,0 +1,90 @@
+
+import json
+import logging
+
+from . request_response_spec import RequestResponse, RequestResponseSpec
+from .. schema import ToolServiceRequest, ToolServiceResponse
+
+logger = logging.getLogger(__name__)
+
+
+class ToolServiceClient(RequestResponse):
+    """Client for invoking dynamically configured tool services."""
+
+    async def call(self, user, config, arguments, timeout=600):
+        """
+        Call a tool service.
+
+        Args:
+            user: User context for multi-tenancy
+            config: Dict of config values (e.g., {"collection": "customers"})
+            arguments: Dict of arguments from LLM
+            timeout: Request timeout in seconds
+
+        Returns:
+            Response string from the tool service
+        """
+        resp = await self.request(
+            ToolServiceRequest(
+                user=user,
+                config=json.dumps(config) if config else "{}",
+                arguments=json.dumps(arguments) if arguments else "{}",
+            ),
+            timeout=timeout
+        )
+
+        if resp.error:
+            raise RuntimeError(resp.error.message)
+
+        return resp.response
+
+    async def call_streaming(self, user, config, arguments, callback, timeout=600):
+        """
+        Call a tool service with streaming response.
+
+        Args:
+            user: User context for multi-tenancy
+            config: Dict of config values
+            arguments: Dict of arguments from LLM
+            callback: Async function called with each response chunk
+            timeout: Request timeout in seconds
+
+        Returns:
+            Final response string
+        """
+        result = []
+
+        async def handle_response(resp):
+            if resp.error:
+                raise RuntimeError(resp.error.message)
+
+            if resp.response:
+                result.append(resp.response)
+                await callback(resp.response)
+
+            return resp.end_of_stream
+
+        await self.request(
+            ToolServiceRequest(
+                user=user,
+                config=json.dumps(config) if config else "{}",
+                arguments=json.dumps(arguments) if arguments else "{}",
+            ),
+            timeout=timeout,
+            recipient=handle_response
+        )
+
+        return "".join(result)
+
+
+class ToolServiceClientSpec(RequestResponseSpec):
+    """Specification for a tool service client."""
+
+    def __init__(self, request_name, response_name):
+        super(ToolServiceClientSpec, self).__init__(
+            request_name=request_name,
+            request_schema=ToolServiceRequest,
+            response_name=response_name,
+            response_schema=ToolServiceResponse,
+            impl=ToolServiceClient,
+        )

trustgraph-base/trustgraph/schema/services/__init__.py

@@ -12,4 +12,5 @@ from .structured_query import *
 from .rows_query import *
 from .diagnosis import *
 from .collection import *
-from .storage import *
+from .storage import *
+from .tool_service import *