trustgraph/trustgraph-cli/trustgraph/cli/set_tool.py

"""
Configures and registers tools in the TrustGraph system.

This script allows you to define agent tools with various types including:
- knowledge-query: Query knowledge bases
- structured-query: Query structured data using natural language
- row-embeddings-query: Semantic search on structured data indexes
- 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.

IMPORTANT: The tool 'name' is used by agents to invoke the tool and must
be a valid function identifier (use snake_case, no spaces or special chars).
The 'description' provides human-readable information about the tool.
"""

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/')
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
default_workspace = os.getenv("TRUSTGRAPH_WORKSPACE", "default")

@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,
        mcp_tool : str,
        collection : str,
        template : str,
        schema_name : str,
        index_name : str,
        limit : int,
        arguments : List[Argument],
        group : List[str],
        state : str,
        applicable_states : List[str],
        token : str = None,
        workspace : str = "default",
):

    api = Api(url, token=token, workspace=workspace).config()

    values = api.get([
        ConfigKey(type="agent", key="tool-index")
    ])

    object = {
        "name": name,
        "description": description,
        "type": type,
    }

    if mcp_tool: object["mcp-tool"] = mcp_tool

    if collection: object["collection"] = collection

    if template: object["template"] = template

    if schema_name: object["schema-name"] = schema_name

    if index_name: object["index-name"] = index_name

    if limit: object["limit"] = limit

    if arguments:
        object["arguments"] = [
            {
                "name": a.name,
                "type": a.type,
                "description": a.description,
            }
            for a in arguments
        ]

    if group is not None: object["group"] = group

    if state: object["state"] = state

    if applicable_states is not None: object["applicable-states"] = applicable_states

    values = api.put([
        ConfigValue(
            type="tool", key=f"{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 (fixed args)
              structured-query     - Query structured data using natural language (fixed args)
              row-embeddings-query - Semantic search on structured data indexes (fixed args)
              text-completion      - Text completion/generation (fixed args)
              mcp-tool             - Model Control Protocol tool (configurable args)
              prompt               - Prompt template query (configurable args)

            Note: Tools marked "(fixed args)" have predefined arguments and don't need
            --argument specified. Tools marked "(configurable args)" require --argument.

            Valid argument types:
              string            - String/text parameter
              number            - Numeric parameter

            Examples:
              %(prog)s --id weather_tool --name get_weather \\
                       --type knowledge-query \\
                       --description "Get weather information for a location" \\
                       --collection weather_data

              %(prog)s --id data_query_tool --name query_data \\
                       --type structured-query \\
                       --description "Query structured data using natural language" \\
                       --collection sales_data

              %(prog)s --id customer_search --name find_customer \\
                       --type row-embeddings-query \\
                       --description "Find customers by name using semantic search" \\
                       --schema-name customers --collection sales \\
                       --index-name full_name --limit 20

              %(prog)s --id calc_tool --name calculate --type mcp-tool \\
                       --description "Perform mathematical calculations" \\
                       --mcp-tool calculator \\
                       --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(
        '-t', '--token',
        default=default_token,
        help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
    )

    parser.add_argument(
        '-w', '--workspace',
        default=default_workspace,
        help=f'Workspace (default: {default_workspace})',
    )

    parser.add_argument(
        '--id',
        help=f'Unique tool identifier',
    )

    parser.add_argument(
        '--name',
        help=f'Tool name used by agents to invoke this tool (use snake_case, e.g., get_weather)',
    )

    parser.add_argument(
        '--description',
        help=f'Detailed description of what the tool does',
    )

    parser.add_argument(
        '--type',
        help=f'Tool type, one of: knowledge-query, structured-query, row-embeddings-query, text-completion, mcp-tool, prompt',
    )

    parser.add_argument(
        '--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, structured-query, and row-embeddings-query types: collection to query',
    )

    parser.add_argument(
        '--schema-name',
        help=f'For row-embeddings-query type: schema name to search within (required)',
    )

    parser.add_argument(
        '--index-name',
        help=f'For row-embeddings-query type: specific index to filter search (optional)',
    )

    parser.add_argument(
        '--limit',
        type=int,
        help=f'For row-embeddings-query type: maximum results to return (default: 10)',
    )

    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)',
    )

    parser.add_argument(
        '--group',
        nargs="*",
        help=f'Tool groups (e.g., read-only, knowledge, admin)',
    )

    parser.add_argument(
        '--state',
        help=f'State to transition to after successful execution',
    )

    parser.add_argument(
        '--applicable-states',
        nargs="*", 
        help=f'States in which this tool is available',
    )

    args = parser.parse_args()

    try:

        valid_types = [
            "knowledge-query", "structured-query", "row-embeddings-query",
            "text-completion", "mcp-tool", "prompt"
        ]

        if args.id is None:
            raise RuntimeError("Must specify --id for tool")

        if args.name is None:
            raise RuntimeError("Must specify --name for tool")

        if args.type:
            if args.type not in valid_types:
                raise RuntimeError(
                    "Type must be one of: " + ", ".join(valid_types)
                )

        mcp_tool = args.mcp_tool

        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,
            mcp_tool=mcp_tool,
            collection=args.collection,
            template=args.template,
            schema_name=args.schema_name,
            index_name=args.index_name,
            limit=args.limit,
            arguments=arguments,
            group=args.group,
            state=args.state,
            applicable_states=args.applicable_states,
            token=args.token,

            workspace=args.workspace,
        )

    except Exception as e:

        print("Exception:", e, flush=True)

if __name__ == "__main__":
    main()