Updated tool CLI to modify tool config

This commit is contained in:
Cyber MacGeddon 2025-07-08 14:43:18 +01:00
parent 15d3d0e961
commit e6bb52aeba
3 changed files with 172 additions and 0 deletions

View file

@ -0,0 +1,170 @@
#!/usr/bin/env python3
"""
Sets a prompt template.
"""
from typing import List
import argparse
import os
from trustgraph.api import Api, ConfigKey, ConfigValue
import json
import tabulate
import textwrap
import dataclasses
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
@dataclasses.dataclass
class Argument:
name : str
type : str
description : str
@staticmethod
def parse(s):
parts = s.split(":")
if len(parts) != 3:
raise RuntimeError(
"Arguments should be form name:type:description"
)
valid_types = [
"string", "number",
]
if parts[1] not in valid_types:
raise RuntimeError(
f"Type {parts[1]} invalid, use: " +
", ".join(valid_types)
)
return Argument(name=parts[0], type=parts[1], description=parts[2])
def set_tool(
url : str,
id : str,
name : str,
description : str,
type : str,
arguments : List[Argument],
):
api = Api(url).config()
values = api.get([
ConfigKey(type="agent", key="tool-index")
])
ix = json.loads(values[0].value)
object = {
"id": id,
"name": name,
"description": description,
"type": type,
"arguments": [
{
"name": a.name,
"type": a.type,
"description": a.description,
}
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)
)
])
print("Tool set.")
def main():
parser = argparse.ArgumentParser(
prog='tg-set-tool',
description=__doc__,
)
parser.add_argument(
'-u', '--api-url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'--id',
help=f'Tool ID',
)
parser.add_argument(
'--name',
help=f'Tool name',
)
parser.add_argument(
'--description',
help=f'Tool description',
)
parser.add_argument(
'--type',
help=f'Tool type, one of: knowledge-query, text-completion, mcp-tool',
)
parser.add_argument(
'--argument',
nargs="*",
help=f'Arguments, form: name:type:description',
)
args = parser.parse_args()
try:
valid_types = [
"knowledge-query", "text-completion", "mcp-tool"
]
if args.id is None:
raise RuntimeError("Must specify --id for prompt")
if args.name is None:
raise RuntimeError("Must specify --name for prompt")
if args.type:
if args.type not in valid_types:
raise RuntimeError(
"Type must be one of: " + ", ".join(valid_types)
)
if args.argument:
arguments = [
Argument.parse(a)
for a in args.argument
]
else:
arguments = []
set_tool(
url=args.api_url, id=args.id, name=args.name,
description=args.description,
type=args.type,
arguments=arguments
)
except Exception as e:
print("Exception:", e, flush=True)
main()

View file

@ -37,6 +37,7 @@ def show_config(url):
table.append(("id", data["id"]))
table.append(("name", data["name"]))
table.append(("description", data["description"]))
table.append(("type", data["type"]))
for n, arg in enumerate(data["arguments"]):
table.append((

View file

@ -70,6 +70,7 @@ setuptools.setup(
"scripts/tg-save-doc-embeds",
"scripts/tg-set-prompt",
"scripts/tg-set-token-costs",
"scripts/tg-set-tool",
"scripts/tg-show-config",
"scripts/tg-show-flow-classes",
"scripts/tg-show-flow-state",