Tool service implementation

This commit is contained in:
Cyber MacGeddon 2026-03-04 13:40:36 +00:00
parent fe3977b5d1
commit 311c4de970
5 changed files with 271 additions and 2 deletions

View file

@ -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"

View file

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