Updated CLI invocation and config model for tools and mcp (#438)

* Updated CLI invocation and config model for tools and mcp

* CLI anomalies

* Tweaked the MCP tool implementation for new model

* Update agent implementation to match the new model

* Fix agent tools, now all tested

* Fixed integration tests

* Fix MCP delete tool params
This commit is contained in:
cybermaggedon 2025-07-16 23:09:32 +01:00 committed by GitHub
parent a96d02da5d
commit 81c7c1181b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 270 additions and 183 deletions

View file

@ -69,39 +69,39 @@ class TestAgentManagerIntegration:
"knowledge_query": Tool(
name="knowledge_query",
description="Query the knowledge graph for information",
arguments={
"question": Argument(
arguments=[
Argument(
name="question",
type="string",
description="The question to ask the knowledge graph"
)
},
],
implementation=KnowledgeQueryImpl,
config={}
),
"text_completion": Tool(
name="text_completion",
description="Generate text completion using LLM",
arguments={
"question": Argument(
arguments=[
Argument(
name="question",
type="string",
description="The question to ask the LLM"
)
},
],
implementation=TextCompletionImpl,
config={}
),
"web_search": Tool(
name="web_search",
description="Search the web for information",
arguments={
"query": Argument(
arguments=[
Argument(
name="query",
type="string",
description="The search query"
)
},
],
implementation=lambda context: AsyncMock(invoke=AsyncMock(return_value="Web search results")),
config={}
)

View file

@ -2,7 +2,7 @@
"""
Deletes MCP (Model Control Protocol) tools from the TrustGraph system.
Removes MCP tool configurations by name from the 'mcp' configuration group.
Removes MCP tool configurations by ID from the 'mcp' configuration group.
"""
import argparse
@ -14,7 +14,7 @@ default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
def delete_mcp_tool(
url : str,
name : str,
id : str,
):
api = Api(url).config()
@ -22,28 +22,28 @@ def delete_mcp_tool(
# Check if the tool exists first
try:
values = api.get([
ConfigKey(type="mcp", key=name)
ConfigKey(type="mcp", key=id)
])
if not values or not values[0].value:
print(f"MCP tool '{name}' not found.")
print(f"MCP tool '{id}' not found.")
return False
except Exception as e:
print(f"MCP tool '{name}' not found.")
print(f"MCP tool '{id}' not found.")
return False
# Delete the MCP tool configuration from the 'mcp' group
try:
api.delete([
ConfigKey(type="mcp", key=name)
ConfigKey(type="mcp", key=id)
])
print(f"MCP tool '{name}' deleted successfully.")
print(f"MCP tool '{id}' deleted successfully.")
return True
except Exception as e:
print(f"Error deleting MCP tool '{name}': {e}")
print(f"Error deleting MCP tool '{id}': {e}")
return False
def main():
@ -56,9 +56,9 @@ def main():
Once deleted, the tool will no longer be available for use.
Examples:
%(prog)s --name weather
%(prog)s --name calculator
%(prog)s --api-url http://localhost:9000/ --name file-reader
%(prog)s --id weather
%(prog)s --id calculator
%(prog)s --api-url http://localhost:9000/ --id file-reader
''').strip(),
formatter_class=argparse.RawDescriptionHelpFormatter
)
@ -70,21 +70,21 @@ def main():
)
parser.add_argument(
'--name',
'--id',
required=True,
help='MCP tool name to delete',
help='MCP tool ID to delete',
)
args = parser.parse_args()
try:
if not args.name:
raise RuntimeError("Must specify --name for MCP tool to delete")
if not args.id:
raise RuntimeError("Must specify --id for MCP tool to delete")
delete_mcp_tool(
url=args.api_url,
name=args.name
id=args.id
)
except Exception as e:

View file

@ -21,27 +21,10 @@ def delete_tool(
api = Api(url).config()
# Get the current tool index
try:
values = api.get([
ConfigKey(type="agent", key="tool-index")
])
ix = json.loads(values[0].value)
except Exception as e:
print(f"Error reading tool index: {e}")
return False
# Check if the tool exists in the index
if id not in ix:
print(f"Tool '{id}' not found in tool index.")
return False
# Check if the tool configuration exists
try:
tool_values = api.get([
ConfigKey(type="agent", key=f"tool.{id}")
ConfigKey(type="tool", key=id)
])
if not tool_values or not tool_values[0].value:
@ -52,22 +35,12 @@ def delete_tool(
print(f"Tool configuration for '{id}' not found.")
return False
# Remove the tool ID from the index
ix.remove(id)
# Delete the tool configuration and update the index
try:
# Update the tool index
api.put([
ConfigValue(
type="agent", key="tool-index", value=json.dumps(ix)
)
])
# Delete the tool configuration
api.delete([
ConfigKey(type="agent", key=f"tool.{id}")
ConfigKey(type="tool", key=id)
])
print(f"Tool '{id}' deleted successfully.")

52
trustgraph-cli/scripts/tg-set-mcp-tool Normal file → Executable file
View file

@ -1,10 +1,17 @@
#!/usr/bin/env python3
"""
Configures and registers MCP (Model Control Protocol) tools in the
TrustGraph system. Allows defining MCP tool configurations with name and
URL. Tools are stored in the 'mcp' configuration group for discovery and
execution.
Configures and registers MCP (Model Context Protocol) tools in the
TrustGraph system.
MCP tools are external services that follow the Model Context Protocol
specification. This script stores MCP tool configurations with:
- id: Unique identifier for the tool
- remote-name: Name used by the MCP server (defaults to id)
- url: MCP server endpoint URL
Configurations are stored in the 'mcp' configuration group and can be
referenced by agent tools using the 'mcp-tool' type.
"""
import argparse
@ -17,7 +24,8 @@ default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
def set_mcp_tool(
url : str,
name : str,
id : str,
remote_name : str,
tool_url : str,
):
@ -26,15 +34,13 @@ def set_mcp_tool(
# Store the MCP tool configuration in the 'mcp' group
values = api.put([
ConfigValue(
type="mcp", key=name, value=json.dumps({
"name": name,
type="mcp", key=id, value=json.dumps({
"remote-name": remote_name,
"url": tool_url,
})
)
])
print(f"MCP tool '{name}' set with URL: {tool_url}")
def main():
parser = argparse.ArgumentParser(
@ -45,8 +51,8 @@ def main():
to the MCP server endpoint that provides the tool functionality.
Examples:
%(prog)s --name weather --tool-url "http://localhost:3000/weather"
%(prog)s --name calculator --tool-url "http://mcp-tools.example.com/calc"
%(prog)s --id weather --tool-url "http://localhost:3000/weather"
%(prog)s --id calculator --tool-url "http://mcp-tools.example.com/calc"
''').strip(),
formatter_class=argparse.RawDescriptionHelpFormatter
)
@ -58,9 +64,15 @@ def main():
)
parser.add_argument(
'--name',
'-i', '--id',
required=True,
help='MCP tool name',
help='MCP tool identifier',
)
parser.add_argument(
'-r', '--remote-name',
required=False,
help='Remote MCP tool name (defaults to --id if not specified)',
)
parser.add_argument(
@ -73,15 +85,21 @@ def main():
try:
if not args.name:
raise RuntimeError("Must specify --name for MCP tool")
if not args.id:
raise RuntimeError("Must specify --id for MCP tool")
if not args.tool_url:
raise RuntimeError("Must specify --url for MCP tool")
raise RuntimeError("Must specify --tool-url for MCP tool")
if args.remote_name:
remote_name = args.remote_name
else:
remote_name = args.id
set_mcp_tool(
url=args.api_url,
name=args.name,
id=args.id,
remote_name=remote_name,
tool_url=args.tool_url
)

View file

@ -2,9 +2,15 @@
"""
Configures and registers tools in the TrustGraph system.
Allows defining tool metadata including ID, name, description, type,
and argument specifications. Tools are stored in the agent configuration
and indexed for discovery and execution.
This script allows you to define agent tools with various types including:
- knowledge-query: Query knowledge bases
- text-completion: Text generation
- mcp-tool: Reference to MCP (Model Context Protocol) tools
- prompt: Prompt template execution
Tools are stored in the 'tool' configuration group and can include
argument specifications for parameterized execution.
"""
from typing import List
@ -51,6 +57,9 @@ def set_tool(
name : str,
description : str,
type : str,
mcp_tool : str,
collection : str,
template : str,
arguments : List[Argument],
):
@ -60,14 +69,20 @@ def set_tool(
ConfigKey(type="agent", key="tool-index")
])
ix = json.loads(values[0].value)
object = {
"id": id,
"name": name,
"description": description,
"type": type,
"arguments": [
}
if mcp_tool: object["mcp-tool"] = mcp_tool
if collection: object["collection"] = collection
if template: object["template"] = template
if arguments:
object["arguments"] = [
{
"name": a.name,
"type": a.type,
@ -75,17 +90,10 @@ def set_tool(
}
for a in arguments
]
}
if id not in ix:
ix.append(id)
values = api.put([
ConfigValue(
type="agent", key="tool-index", value=json.dumps(ix)
),
ConfigValue(
type="agent", key=f"tool.{id}", value=json.dumps(object)
type="tool", key=f"{id}", value=json.dumps(object)
)
])
@ -100,7 +108,8 @@ def main():
Valid tool types:
knowledge-query - Query knowledge bases
text-completion - Text completion/generation
mcp-tool - Model Control Protocol tool
mcp-tool - Model Control Protocol tool
prompt - Prompt template query
Valid argument types:
string - String/text parameter
@ -128,28 +137,43 @@ def main():
parser.add_argument(
'--id',
help=f'Tool ID',
help=f'Unique tool identifier',
)
parser.add_argument(
'--name',
help=f'Tool name',
help=f'Human-readable tool name',
)
parser.add_argument(
'--description',
help=f'Tool description',
help=f'Detailed description of what the tool does',
)
parser.add_argument(
'--type',
help=f'Tool type, one of: knowledge-query, text-completion, mcp-tool',
help=f'Tool type, one of: knowledge-query, text-completion, mcp-tool, prompt',
)
parser.add_argument(
'--argument',
nargs="*",
help=f'Arguments, form: name:type:description',
'--mcp-tool',
help=f'For MCP type: ID of MCP tool configuration (as defined by tg-set-mcp-tool)',
)
parser.add_argument(
'--collection',
help=f'For knowledge-query type: collection to query',
)
parser.add_argument(
'--template',
help=f'For prompt type: template ID to use',
)
parser.add_argument(
'--argument',
nargs="*",
help=f'Tool arguments in the form: name:type:description (can specify multiple)',
)
args = parser.parse_args()
@ -157,14 +181,14 @@ def main():
try:
valid_types = [
"knowledge-query", "text-completion", "mcp-tool"
"knowledge-query", "text-completion", "mcp-tool", "prompt"
]
if args.id is None:
raise RuntimeError("Must specify --id for prompt")
raise RuntimeError("Must specify --id for tool")
if args.name is None:
raise RuntimeError("Must specify --name for prompt")
raise RuntimeError("Must specify --name for tool")
if args.type:
if args.type not in valid_types:
@ -172,6 +196,8 @@ def main():
"Type must be one of: " + ", ".join(valid_types)
)
mcp_tool = args.mcp_tool
if args.argument:
arguments = [
Argument.parse(a)
@ -181,10 +207,15 @@ def main():
arguments = []
set_tool(
url=args.api_url, id=args.id, name=args.name,
url=args.api_url,
id=args.id,
name=args.name,
description=args.description,
type=args.type,
arguments=arguments
mcp_tool=mcp_tool,
collection=args.collection,
template=args.template,
arguments=arguments,
)
except Exception as e:

View file

@ -1,7 +1,7 @@
#!/usr/bin/env python3
"""
Dumps out the current agent tool configuration
Displays the current MCP (Model Context Protocol) tool configuration
"""
import argparse
@ -26,11 +26,10 @@ def show_config(url):
table = []
table.append(("id", value.key))
table.append(("name", data["name"]))
table.append(("remote-name", data["remote-name"]))
table.append(("url", data["url"]))
print()
print(value.key + ":")
print(tabulate.tabulate(
table,

View file

@ -1,7 +1,13 @@
#!/usr/bin/env python3
"""
Dumps out the current agent tool configuration
Displays the current agent tool configurations
Shows all configured tools including their types:
- knowledge-query: Tools that query knowledge bases
- text-completion: Tools for text generation
- mcp-tool: References to MCP (Model Context Protocol) tools
- prompt: Tools that execute prompt templates
"""
import argparse
@ -17,37 +23,37 @@ def show_config(url):
api = Api(url).config()
values = api.get([
ConfigKey(type="agent", key="tool-index")
])
values = api.get_values(type="tool")
ix = json.loads(values[0].value)
for item in values:
values = api.get([
ConfigKey(type="agent", key=f"tool.{v}")
for v in ix
])
id = item.key
data = json.loads(item.value)
for n, key in enumerate(ix):
data = json.loads(values[n].value)
tp = data["type"]
table = []
table.append(("id", data["id"]))
table.append(("id", id))
table.append(("name", data["name"]))
table.append(("description", data["description"]))
table.append(("type", data["type"]))
table.append(("type", tp))
for n, arg in enumerate(data["arguments"]):
table.append((
f"arg {n}",
f"{arg['name']}: {arg['type']}\n{arg['description']}"
))
if tp == "mcp-tool":
table.append(("mcp-tool", data["mcp-tool"]))
if tp == "knowledge-query":
table.append(("collection", data["collection"]))
if tp == "prompt":
table.append(("template", data["template"]))
for n, arg in enumerate(data["arguments"]):
table.append((
f"arg {n}",
f"{arg['name']}: {arg['type']}\n{arg['description']}"
))
print()
print(key + ":")
print(tabulate.tabulate(
table,

View file

@ -47,8 +47,8 @@ class Service(ToolService):
url = self.mcp_services[name]["url"]
if "name" in self.mcp_services[name]:
remote_name = self.mcp_services[name]["name"]
if "remote-name" in self.mcp_services[name]:
remote_name = self.mcp_services[name]["remote-name"]
else:
remote_name = name

View file

@ -39,7 +39,7 @@ class AgentManager:
"type": arg.type,
"description": arg.description
}
for arg in tool.arguments.values()
for arg in tool.arguments
]
}
for tool in self.tools.values()

View file

@ -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:
# 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)
pc = config[f"tool.{k}"]
data = json.loads(pc)
impl_id = data.get("type")
name = data.get("name")
arguments = {
v.get("name"): Argument(
name = v.get("name"),
type = v.get("type"),
description = v.get("description")
)
for v in data["arguments"]
}
# 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"
)
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[name] = Tool(
name=name,
description=data.get("description"),
implementation=impl,
config=data, # Store full config for reference
arguments=arguments,
)
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:

View file

@ -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)
@ -51,3 +81,21 @@ class McpToolImpl:
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.prompt(
id=self.template_id,
variables=arguments
)