diff --git a/trustgraph-cli/scripts/tg-set-tool b/trustgraph-cli/scripts/tg-set-tool new file mode 100755 index 00000000..479a4315 --- /dev/null +++ b/trustgraph-cli/scripts/tg-set-tool @@ -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() + diff --git a/trustgraph-cli/scripts/tg-show-tools b/trustgraph-cli/scripts/tg-show-tools index b6c4a8e4..2056a520 100755 --- a/trustgraph-cli/scripts/tg-show-tools +++ b/trustgraph-cli/scripts/tg-show-tools @@ -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(( diff --git a/trustgraph-cli/setup.py b/trustgraph-cli/setup.py index 6af3f38e..c9afe88e 100644 --- a/trustgraph-cli/setup.py +++ b/trustgraph-cli/setup.py @@ -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",