trustgraph/trustgraph-cli/scripts/tg-set-tool

#!/usr/bin/env python3

"""
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.
"""

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__,
        epilog=textwrap.dedent('''
            Valid tool types:
              knowledge-query    - Query knowledge bases
              text-completion    - Text completion/generation
              mcp-tool          - Model Control Protocol tool
            
            Valid argument types:
              string            - String/text parameter
              number            - Numeric parameter
            
            Examples:
              %(prog)s --id weather --name "Weather lookup" \\
                       --type knowledge-query \\
                       --description "Get weather information" \\
                       --argument location:string:"Location to query" \\
                       --argument units:string:"Temperature units (C/F)"
              
              %(prog)s --id calculator --name "Calculator" --type mcp-tool \\
                       --description "Perform calculations" \\
                       --argument expression:string:"Mathematical expression"
        ''').strip(),
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    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()