🛠️ Creating Custom Tools

Custom tools are reusable components that encapsulate specific functionality in your Loopstack workflows. They can accept arguments, use dependency injection, and integrate seamlessly with the workflow execution.

Basic Tool Structure

A custom tool implements the ToolInterface and the execute method:


import { Inject } from '@nestjs/common';
import { z } from 'zod';
import { Input, Tool, ToolInterface, ToolResult } from '@loopstack/common';
import { MathService } from '../services/math.service';
 
export type MathSumArgs = {
  a: number;
  b: number;
};
 
@Tool({
  config: {
    description: 'Math tool calculating the sum of two arguments by using an injected service.',
  },
})
export class MathSumTool implements ToolInterface {
  @Inject()
  private mathService: MathService;
 
  @Input({
    schema: z
      .object({
        a: z.number(),
        b: z.number(),
      })
      .strict(),
  })
  args: MathSumArgs;
 
  async execute(args: MathSumArgs): Promise<ToolResult<number>> {
    const sum = this.mathService.sum(args.a, args.b);
    return Promise.resolve({
      data: sum,
    });
  }
}

Key Components

1. Required Decorators

@Tool

Defines the tool’s metadata and configuration
config.description: Human-readable description of what the tool does

@Input({ schema })

Defines the expected arguments using a Zod schema
Applied to a class property that holds the arguments

2. Arguments Schema

Define the expected arguments using @Input with a Zod schema on a class property, and export a type for use in the execute method:


export type MathSumArgs = {
  a: number;
  b: number;
};
 
@Input({
  schema: z
    .object({
      a: z.number(),
      b: z.number(),
    })
    .strict(),
})
args: MathSumArgs;

3. Base Class

Implements ToolInterface:


export class MathSumTool implements ToolInterface {
  // ...
}

4. Execute Method

The execute() method contains your tool’s logic. Arguments are passed directly as a parameter:


async execute(args: MathSumArgs): Promise<ToolResult<number>> {
  const sum = this.mathService.sum(args.a, args.b);
  return Promise.resolve({
    data: sum,
  });
}

5. Dependency Injection

Tools support NestJS dependency injection. Use @Inject() to inject services:


@Inject()
private mathService: MathService;

Stateful Tools

Tools can maintain internal state across invocations. For example, a counter tool that increments each time it is called:


import { Tool, ToolInterface, ToolResult } from '@loopstack/common';
 
@Tool({
  config: {
    description: 'Counter tool.',
  },
})
export class CounterTool implements ToolInterface {
  count: number = 0;
 
  async execute(): Promise<ToolResult> {
    this.count++;
    return Promise.resolve({
      data: this.count,
    });
  }
}

This tool requires no arguments and returns an incrementing count each time it is executed.

Using Tools in Workflows

1. Define the Workflow Class

Use the @InjectTool() decorator to register tools in your workflow class. You can also define input arguments, state, output, and helper functions using their respective decorators:


import { z } from 'zod';
import { DefineHelper, InjectTool, Input, Output, State, Workflow } from '@loopstack/common';
import { CreateChatMessage } from '@loopstack/create-chat-message-tool';
import { MathSumTool } from '../tools';
import { CounterTool } from '../tools';
 
@Workflow({
  configFile: __dirname + '/custom-tool-example.workflow.yaml',
})
export class CustomToolExampleWorkflow {
  @InjectTool() private counterTool: CounterTool;
  @InjectTool() private createChatMessage: CreateChatMessage;
  @InjectTool() private mathTool: MathSumTool;
 
  @Input({
    schema: z
      .object({
        a: z.number().default(1),
        b: z.number().default(2),
      })
      .strict(),
  })
  args: {
    a: number;
    b: number;
  };
 
  @State({
    schema: z
      .object({
        total: z.number().optional(),
        count1: z.number().optional(),
        count2: z.number().optional(),
        count3: z.number().optional(),
      })
      .strict(),
  })
  state: {
    total?: number;
    count1?: number;
    count2?: number;
    count3?: number;
  };
 
  @Output()
  result() {
    return {
      total: this.state.total,
    };
  }
 
  @DefineHelper()
  sum(a: number, b: number) {
    return a + b;
  }
}

@Input({ schema }) defines the workflow’s input arguments with a Zod schema and default values.

@State({ schema }) defines the workflow’s mutable state that can be written to during transitions via assign.

@Output() defines a method that returns the workflow’s output when it completes.

@DefineHelper() defines helper functions that can be called from workflow YAML templates (e.g. {{ sum args.a args.b }}).

2. Register in a Module

Add the workflow, tools, and their dependencies to your module. Use imports for tool modules that provide their own providers:


import { Module } from '@nestjs/common';
import { CreateChatMessageToolModule } from '@loopstack/create-chat-message-tool';
import { MathService } from './services/math.service';
import { MathSumTool } from './tools';
import { CounterTool } from './tools';
import { CustomToolExampleWorkflow } from './workflows';
 
@Module({
  imports: [CreateChatMessageToolModule],
  providers: [CustomToolExampleWorkflow, MathSumTool, CounterTool, MathService],
  exports: [CustomToolExampleWorkflow, MathSumTool, CounterTool, MathService],
})
export class CustomToolModule {}

3. Call in Workflow YAML

Use the tool in your workflow transitions. Reference tools by their property name (camelCase):


title: 'Custom Tool'
 
description: |
  This workflow demonstrates the usage of custom tools, including both stateless and stateful tools.
  It performs a simple addition operation using a custom MathSumTool and showcases the behavior of
  stateless and stateful counter tools.
 
ui:
  form:
    properties:
      a:
        title: 'First number (a)'
      b:
        title: 'Second number (b)'
 
transitions:
  - id: calculate
    from: start
    to: end
    call:
      # Use a custom tool
      - id: calculation
        tool: mathTool
        args:
          a: ${{ args.a }}
          b: ${{ args.b }}
        assign:
          total: ${{ result.data }}
 
      # Display the result
      - tool: createChatMessage
        args:
          role: 'assistant'
          content: |
            Tool calculation result:
            {{ args.a }} + {{ args.b }} = {{ state.total }}
 
      # Alternatively, use a custom workflow helper function
      - tool: createChatMessage
        args:
          role: 'assistant'
          content: |
            Alternatively, using workflow helper function:
            {{ args.a }} + {{ args.b }} = {{ sum args.a args.b }}
 
      # Use a stateful tool
      - id: count1
        tool: counterTool
        assign:
          count1: ${{ result.data }}
      - id: count2
        tool: counterTool
        assign:
          count2: ${{ result.data }}
      - id: count3
        tool: counterTool
        assign:
          count3: ${{ result.data }}
      - tool: createChatMessage
        args:
          role: 'assistant'
          content: |
            Counter tool should count:
 
            {{ state.count1 }}, {{ state.count2 }}, {{ state.count3 }}