From 311c4de970457d7dcbb7d42e49c3c8c72ed0c77d Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 4 Mar 2026 13:40:36 +0000 Subject: [PATCH] Tool service implementation --- trustgraph-base/trustgraph/base/__init__.py | 1 + .../trustgraph/base/tool_service_client.py | 90 ++++++++++++++ .../trustgraph/schema/services/__init__.py | 3 +- .../trustgraph/agent/react/service.py | 66 +++++++++- .../trustgraph/agent/react/tools.py | 113 ++++++++++++++++++ 5 files changed, 271 insertions(+), 2 deletions(-) create mode 100644 trustgraph-base/trustgraph/base/tool_service_client.py diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index 557109a2..795c5c56 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/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 diff --git a/trustgraph-base/trustgraph/base/tool_service_client.py b/trustgraph-base/trustgraph/base/tool_service_client.py new file mode 100644 index 00000000..81930ba0 --- /dev/null +++ b/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, + ) diff --git a/trustgraph-base/trustgraph/schema/services/__init__.py b/trustgraph-base/trustgraph/schema/services/__init__.py index 7b40ca0a..f246bc31 100644 --- a/trustgraph-base/trustgraph/schema/services/__init__.py +++ b/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 * \ No newline at end of file +from .storage import * +from .tool_service import * \ 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 1a44ef9e..c4cb1fab 100755 --- a/trustgraph-flow/trustgraph/agent/react/service.py +++ b/trustgraph-flow/trustgraph/agent/react/service.py @@ -17,7 +17,7 @@ from ... base import RowEmbeddingsQueryClientSpec, EmbeddingsClientSpec from ... schema import AgentRequest, AgentResponse, AgentStep, Error -from . tools import KnowledgeQueryImpl, TextCompletionImpl, McpToolImpl, PromptImpl, StructuredQueryImpl, RowEmbeddingsQueryImpl +from . tools import KnowledgeQueryImpl, TextCompletionImpl, McpToolImpl, PromptImpl, StructuredQueryImpl, RowEmbeddingsQueryImpl, ToolServiceImpl from . agent_manager import AgentManager from ..tool_filter import validate_tool_config, filter_tools_by_group_and_state, get_next_state @@ -51,6 +51,9 @@ class Processor(AgentService): additional_context="", ) + # Track active tool service clients for cleanup + self.tool_service_clients = {} + self.config_handlers.append(self.on_tools_config) self.register_specification( @@ -110,6 +113,16 @@ class Processor(AgentService): tools = {} + # Load tool-service configurations first + tool_services = {} + if "tool-service" in config: + for service_id, service_value in config["tool-service"].items(): + service_data = json.loads(service_value) + tool_services[service_id] = service_data + logger.debug(f"Loaded tool-service config: {service_id}") + + logger.info(f"Loaded {len(tool_services)} tool-service configurations") + # Load tool configurations from the new location if "tool" in config: for tool_id, tool_value in config["tool"].items(): @@ -177,6 +190,57 @@ class Processor(AgentService): limit=int(data.get("limit", 10)) # Max results ) arguments = RowEmbeddingsQueryImpl.get_arguments() + elif impl_id == "tool-service": + # Dynamic tool service - look up the service config + service_ref = data.get("service") + if not service_ref: + raise RuntimeError( + f"Tool {name} has type 'tool-service' but no 'service' reference" + ) + if service_ref not in tool_services: + raise RuntimeError( + f"Tool {name} references unknown tool-service '{service_ref}'" + ) + + service_config = tool_services[service_ref] + service_topic = service_config.get("topic") + if not service_topic: + raise RuntimeError( + f"Tool-service '{service_ref}' has no 'topic' defined" + ) + + # Build config values from tool config + # Extract any config params defined by the service + config_params = service_config.get("config-params", []) + config_values = {} + for param in config_params: + param_name = param.get("name") if isinstance(param, dict) else param + if param_name in data: + config_values[param_name] = data[param_name] + elif isinstance(param, dict) and param.get("required", False): + raise RuntimeError( + f"Tool {name} missing required config param '{param_name}'" + ) + + # Arguments come from tool config + config_args = data.get("arguments", []) + arguments = [ + Argument( + name=arg.get("name"), + type=arg.get("type"), + description=arg.get("description") + ) + for arg in config_args + ] + + # Store topic for the implementation + impl = functools.partial( + ToolServiceImpl, + service_topic=service_topic, + config_values=config_values, + arguments=arguments, + processor=self, + ) 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 2b442a0d..7bfa67ea 100644 --- a/trustgraph-flow/trustgraph/agent/react/tools.py +++ b/trustgraph-flow/trustgraph/agent/react/tools.py @@ -202,3 +202,116 @@ class PromptImpl: id=self.template_id, variables=arguments ) + + +# This tool implementation invokes a dynamically configured tool service +class ToolServiceImpl: + """ + Implementation for dynamically pluggable tool services. + + Tool services are external Pulsar services that can be invoked as agent tools. + The service is configured via a tool-service descriptor that defines the topic, + and a tool descriptor that provides config values and argument definitions. + """ + + def __init__(self, context, service_topic, config_values=None, arguments=None, processor=None): + """ + Initialize a tool service implementation. + + Args: + context: The context function (provides user info) + service_topic: The base topic name for the service + config_values: Dict of config values (e.g., {"collection": "customers"}) + arguments: List of Argument objects defining the tool's parameters + processor: The Processor instance (for pubsub access) + """ + self.context = context + self.service_topic = service_topic + self.config_values = config_values or {} + self.arguments = arguments or [] + self.processor = processor + self._client = None + + def get_arguments(self): + return self.arguments + + async def _get_or_create_client(self): + """Get or create the tool service client.""" + if self._client is not None: + return self._client + + # Check if processor already has a client for this topic + if self.service_topic in self.processor.tool_service_clients: + self._client = self.processor.tool_service_clients[self.service_topic] + return self._client + + # Import here to avoid circular imports + from ....base.tool_service_client import ToolServiceClient + from ....base.metrics import ProducerMetrics, SubscriberMetrics + from ....schema import ToolServiceRequest, ToolServiceResponse + import uuid + + request_topic = f"{self.service_topic}-request" + response_topic = f"{self.service_topic}-response" + + request_metrics = ProducerMetrics( + processor=self.processor.id, + flow="tool-service", + name=request_topic + ) + response_metrics = SubscriberMetrics( + processor=self.processor.id, + flow="tool-service", + name=response_topic + ) + + # Create unique subscription for responses + subscription = f"{self.processor.id}--tool-service--{self.service_topic}--{uuid.uuid4()}" + + self._client = ToolServiceClient( + backend=self.processor.pubsub, + subscription=subscription, + consumer_name=self.processor.id, + request_topic=request_topic, + request_schema=ToolServiceRequest, + request_metrics=request_metrics, + response_topic=response_topic, + response_schema=ToolServiceResponse, + response_metrics=response_metrics, + ) + + # Start the client + await self._client.start() + + # Register for cleanup + self.processor.tool_service_clients[self.service_topic] = self._client + + logger.debug(f"Created tool service client for {self.service_topic}") + return self._client + + async def invoke(self, **arguments): + logger.debug(f"Tool service invocation: {self.service_topic}...") + logger.debug(f"Config: {self.config_values}") + logger.debug(f"Arguments: {arguments}") + + # Get user from context if available + user = "trustgraph" + if hasattr(self.context, '_user'): + user = self.context._user + + # Get or create the client + client = await self._get_or_create_client() + + # Call the tool service + response = await client.call( + user=user, + config=self.config_values, + arguments=arguments, + ) + + logger.debug(f"Tool service response: {response}") + + if isinstance(response, str): + return response + else: + return json.dumps(response)