diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py index b28be1a6..3e4dfe64 100755 --- a/trustgraph-flow/trustgraph/agent/react/service.py +++ b/trustgraph-flow/trustgraph/agent/react/service.py @@ -12,7 +12,7 @@ from ... base import GraphRagClientSpec, ToolClientSpec from ... schema import AgentRequest, AgentResponse, AgentStep, Error -from . tools import KnowledgeQueryImpl, TextCompletionImpl, McpToolImpl +from . tools import KnowledgeQueryImpl, TextCompletionImpl, McpToolImpl, PromptImpl from . agent_manager import AgentManager from . types import Final, Action, Tool, Argument @@ -79,64 +79,76 @@ class Processor(AgentService): print("Loading configuration version", version) - if self.config_key not in config: - print(f"No key {self.config_key} in config", flush=True) - return - - config = config[self.config_key] - try: - # This is some extra stuff to put in the prompt - additional = config.get("additional-context", None) - - ix = json.loads(config["tool-index"]) - tools = {} - for k in ix: - - pc = config[f"tool.{k}"] - data = json.loads(pc) - - arguments = { - v.get("name"): Argument( - name = v.get("name"), - type = v.get("type"), - description = v.get("description") + # Load tool configurations from the new location + if "tool" in config: + for tool_id, tool_value in config["tool"].items(): + data = json.loads(tool_value) + + impl_id = data.get("type") + name = data.get("name") + + # Create the appropriate implementation + if impl_id == "knowledge-query": + impl = functools.partial( + KnowledgeQueryImpl, + collection=data.get("collection") + ) + arguments = KnowledgeQueryImpl.get_arguments() + elif impl_id == "text-completion": + impl = TextCompletionImpl + arguments = TextCompletionImpl.get_arguments() + elif impl_id == "mcp-tool": + impl = functools.partial( + McpToolImpl, + mcp_tool_id=data.get("mcp-tool") + ) + arguments = McpToolImpl.get_arguments() + elif impl_id == "prompt": + # For prompt tools, arguments come from 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 + ] + impl = functools.partial( + PromptImpl, + template_id=data.get("template"), + arguments=arguments + ) + else: + raise RuntimeError( + f"Tool type {impl_id} not known" + ) + + tools[name] = Tool( + name=name, + description=data.get("description"), + implementation=impl, + config=data, # Store full config for reference + arguments=arguments, ) - for v in data["arguments"] - } - - impl_id = data.get("type") - - name = data.get("name") - - if impl_id == "knowledge-query": - impl = KnowledgeQueryImpl - elif impl_id == "text-completion": - impl = TextCompletionImpl - elif impl_id == "mcp-tool": - impl = functools.partial(McpToolImpl, name=k) - else: - raise RuntimeError( - f"Tool-kind {impl_id} not known" - ) - - tools[data.get("name")] = Tool( - name = name, - description = data.get("description"), - implementation = impl, - config=data.get("config", {}), - arguments = arguments, - ) - + + # Load additional context from agent config if it exists + additional = None + if self.config_key in config: + agent_config = config[self.config_key] + additional = agent_config.get("additional-context", None) + self.agent = AgentManager( tools=tools, additional_context=additional ) - print("Prompt configuration reloaded.", flush=True) + print(f"Loaded {len(tools)} tools", flush=True) + print("Tool configuration reloaded.", flush=True) except Exception as e: diff --git a/trustgraph-flow/trustgraph/agent/react/tools.py b/trustgraph-flow/trustgraph/agent/react/tools.py index a4ba9907..09c6cd92 100644 --- a/trustgraph-flow/trustgraph/agent/react/tools.py +++ b/trustgraph-flow/trustgraph/agent/react/tools.py @@ -1,11 +1,24 @@ import json +from .types import Argument # This tool implementation knows how to put a question to the graph RAG # service class KnowledgeQueryImpl: - def __init__(self, context): + def __init__(self, context, collection=None): self.context = context + self.collection = collection + + @staticmethod + def get_arguments(): + return [ + Argument( + name="question", + type="string", + description="The question to ask the knowledge base" + ) + ] + async def invoke(self, **arguments): client = self.context("graph-rag-request") print("Graph RAG question...", flush=True) @@ -18,6 +31,17 @@ class KnowledgeQueryImpl: class TextCompletionImpl: def __init__(self, context): self.context = context + + @staticmethod + def get_arguments(): + return [ + Argument( + name="question", + type="string", + description="The text prompt or question for completion" + ) + ] + async def invoke(self, **arguments): client = self.context("prompt-request") print("Prompt question...", flush=True) @@ -29,18 +53,24 @@ class TextCompletionImpl: # the mcp-tool service. class McpToolImpl: - def __init__(self, context, name): + def __init__(self, context, mcp_tool_id): self.context = context - self.name = name + self.mcp_tool_id = mcp_tool_id + + @staticmethod + def get_arguments(): + # MCP tools define their own arguments dynamically + # For now, we return empty list and let the MCP service handle validation + return [] async def invoke(self, **arguments): client = self.context("mcp-tool-request") - print(f"MCP tool invocation: {self.name}...", flush=True) + print(f"MCP tool invocation: {self.mcp_tool_id}...", flush=True) output = await client.invoke( - name = self.name, - parameters = {}, + name = self.mcp_tool_id, + parameters = arguments, # Pass the actual arguments ) print(output) @@ -50,4 +80,22 @@ class McpToolImpl: else: return json.dumps(output) - + +# This tool implementation knows how to execute prompt templates +class PromptImpl: + def __init__(self, context, template_id, arguments=None): + self.context = context + self.template_id = template_id + self.arguments = arguments or [] # These come from config + + def get_arguments(self): + # For prompt tools, arguments are defined in configuration + return self.arguments + + async def invoke(self, **arguments): + client = self.context("prompt-request") + print(f"Prompt template invocation: {self.template_id}...", flush=True) + return await client.template( + template_id=self.template_id, + variables=arguments + )