Implementing Tool Functionality in Conversational AI

import { routeMessage } from './routeMessage';
import { expect, it } from 'vitest';
import { z } from 'zod';

it('uses tools if relevant tools are available', async () => {
  const plan = await routeMessage({
    messages: [
      {
        content: 'What is 2+2?',
        role: 'user' as const,
      },
    ],
    tools: [
      {
        description: 'Adds two numbers.',
        execute: async ({ a, b }) => {
          return {
            sum: a + b,
          };
        },
        name: 'addNumbers',
        parameters: z.object({
          a: z.number(),
          b: z.number(),
        }),
        response: z.object({
          sum: z.number(),
        }),
      },
    ],
  });

  expect(plan).toEqual({
    actions: [
      {
        name: 'addNumbers',
        parameters: {
          a: 2,
          b: 2,
        },
      },
    ],
  });
});

import { zodToJsonSchema } from 'zod-to-json-schema';
import { type AnyZodObject, z } from 'zod';

type Tool<P extends AnyZodObject, R extends AnyZodObject> = {
  description: string;
  execute: (parameters?: z.infer<P>) => Promise<z.infer<R>>;
  name: string;
  parameters?: P;
  response: R;
};

const describeTool = (tool: Tool<AnyZodObject, AnyZodObject>) => {
  return {
    description: tool.description,
    name: tool.name,
    parameters: tool.parameters ? zodToJsonSchema(tool.parameters) : null,
  };
};

import { type CoreMessage } from 'ai';

export const routeMessage = async ({
  messages,
  tools,
}: {
  messages: CoreMessage[];
  tools: ReadonlyArray<Tool<AnyZodObject, AnyZodObject>>;
}) => {}

You are an assistant with access to the following tools:

[tools]

Tools are described in the following format:

* "description" describes what the tool does.
* "name" is the name of the tool.
* "parameters" is the JSON schema of the tool parameters (or null if the tool does not have parameters).

You are also given the following conversation history:

[messages]

The conversation history is a list of messages exchanged between the user and the assistant. It may also describe previous actions taken by the assistant.

Based on the conversation history, and the tools you have access to, propose a plan for how to answer the user's question.

The response should be a JSON object with "actions" property, which is an array of tools to use. Each tool is represented as an object with the following properties:

* "name": the name of the tool to use.
* "parameters": the parameters to pass to the tool (or null if the tool does not have parameters).

The same tool can be used multiple times in the plan.

If the conversation does not necessitate the use of tools, respond with an empty action plan, e.g.,

{
  actions: [],
}

import { quickPrompt } from './quickPrompt';
import { type CoreMessage } from 'ai';
import multiline from 'multiline-ts';
import { type AnyZodObject, z } from 'zod';
import { zodToJsonSchema } from 'zod-to-json-schema';

const SerializableZodSchema = z.record(
  z.union([z.string(), z.number(), z.boolean()]),
);

type Tool<P extends AnyZodObject, R extends AnyZodObject> = {
  description: string;
  execute: (parameters?: z.infer<P>) => Promise<z.infer<R>>;
  name: string;
  parameters?: P;
  response: R;
};

const describeTool = (tool: Tool<AnyZodObject, AnyZodObject>) => {
  return {
    description: tool.description,
    name: tool.name,
    parameters: tool.parameters ? zodToJsonSchema(tool.parameters) : null,
  };
};

export const routeMessage = async ({
  messages,
  tools,
}: {
  messages: CoreMessage[];
  tools: ReadonlyArray<Tool<AnyZodObject, AnyZodObject>>;
}) => {
  return await quickPrompt({
    model: 'openai@gpt-4o-mini',
    name: 'routeMessage',
    prompt: multiline`
      You are an assistant with access to the following tools:

      ${JSON.stringify(
        tools.map((tool) => {
          return describeTool(tool);
        })
      )}

      Tools are described in the following format:

      * "description" describes what the tool does.
      * "name" is the name of the tool.
      * "parameters" is the JSON schema of the tool parameters (or null if the tool does not have parameters).

      You are also given the following conversation history:

      ${JSON.stringify(messages)}

      The conversation history is a list of messages exchanged between the user and the assistant. It may also describe previous actions taken by the assistant.

      Based on the conversation history, and the tools you have access to, propose a plan for how to answer the user's question.

      The response should be a JSON object with "actions" property, which is an array of tools to use. Each tool is represented as an object with the following properties:

      * "name": the name of the tool to use.
      * "parameters": the parameters to pass to the tool (or null if the tool does not have parameters).

      The same tool can be used multiple times in the plan.

      If the conversation does not necessitate the use of tools, respond with an empty action plan, e.g.,

      {
        actions: [],
      }
    `,
    zodSchema: z.object({
      actions: z.array(
        z.object({
          name: z.string(),
          parameters: SerializableZodSchema.nullable(),
        }),
      ),
    }),
  });
};

// We may have multiple cycles of invocations.
// See the explanation after this code example.
while (toolCycle < MAX_TOOL_CYCLES) {
  const plan = await routeMessage({
    messages,
    tools,
  });

  // If `routeMessage` returns an empty plan, it means that the conversation does not require the use of tools.
  if (plan.actions.length === 0) {
    break;
  }

  // For each tool we use, we need to record the invocation and the result in the conversation history.
  for (const action of plan.actions) {
    const toolCallId = randomUUID();

    messages.push({
      content: [
        {
          args: action.parameters,
          toolCallId,
          toolName: action.name,
          type: 'tool-call' as const,
        },
      ],
      role: 'assistant',
    });

    const result = await invokeUserTool(action.name, action.parameters);

    messages.push({
      content: [
        {
          result,
          toolCallId,
          toolName: action.name,
          type: 'tool-result' as const,
        },
      ],
      role: 'tool',
    });
  }
}

// Now pass the messages history to the LLM which will use the recorded tool calls to generate a response.
await streamAssistantResponse({
  messages,
  signal: abortController.signal,
  visitor: context.visitor,
});