mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-21 19:21:03 +02:00
170 lines
3.5 KiB
Python
Executable file
170 lines
3.5 KiB
Python
Executable file
#!/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()
|
|
|