In this article, we will review “tool” function in vercel/ai package source code.

Let me provide an explanation as to how I ended up finding this. This below code snippet is picked from create-document.ts

export const createDocument = ({ session, dataStream }: CreateDocumentProps) =>
  tool({
    description:
      'Create a document for a writing or content creation activities. This tool will call other functions that will generate the contents of the document based on the title and kind.',
    parameters: z.object({
      title: z.string(),
      kind: z.enum(artifactKinds),
    }),
    execute: async ({ title, kind }) => {

I couldn’t help but notice the tool that accepts two parameters:

1. an object contains description, parameters

2. execute function

and this tool function is imported at the top of this create-document.ts as shown below:

import { DataStreamWriter, tool } from 'ai';

I am aware that “ai” package is owned by Vercel, don’t ask me how. I just remember it from some OSS code I read some time ago.

So once I found the source code of this “ai” package, I started digging into the ai source code to find tool function definition.

After a while, I found this tool definition in vercel/ai/packages/ai/core/tool/tool.ts and contains about 132 LOC at the time of this writing.

As you can see from the above screenshot, defining types/interface along with comments makes up for so many LOC in this file.

export type inferParameters<PARAMETERS extends ToolParameters> =
      : never;

export interface ToolExecutionOptions {
}

/**
A tool contains the description and the schema of the input that the tool expects.
This enables the language model to generate the input.

The tool can also contain an optional execute function for the actual execution function of the tool.
 */
export type Tool<PARAMETERS extends ToolParameters = any, RESULT = any> = {
);

The actual definition of this tool function looks like below:

/**
Helper function for inferring the execute args of a tool.
 */
// Note: special type inference is needed for the execute function args to make sure they are inferred correctly.
export function tool<PARAMETERS extends ToolParameters, RESULT>(
  tool: Tool<PARAMETERS, RESULT> & {
    execute: (
      args: inferParameters<PARAMETERS>,
      options: ToolExecutionOptions,
    ) => PromiseLike<RESULT>;
  },
): Tool<PARAMETERS, RESULT> & {
  execute: (
    args: inferParameters<PARAMETERS>,
    options: ToolExecutionOptions,
  ) => PromiseLike<RESULT>;
};
export function tool<PARAMETERS extends ToolParameters, RESULT>(
  tool: Tool<PARAMETERS, RESULT> & {
    execute?: undefined;
  },
): Tool<PARAMETERS, RESULT> & {
  execute: undefined;
};
export function tool<PARAMETERS extends ToolParameters, RESULT = any>(
  tool: Tool<PARAMETERS, RESULT>,
): Tool<PARAMETERS, RESULT> {
  return tool;
}

What you are seeing is function overloading in TypeScript, I have seen this way too many times in the OSS codebases. I was wondering what’s the actual implementation of tool function since it’s got function overloading.

I asked ChatGPT and it said that the actual tool implementation is that it just returns tool function itself. Function overloading is for type inferences.

About me:

Hey, my name is Ramu Narasinga. I study large open-source projects and create content about their codebase architecture and best practices, sharing it through articles, videos.

I am open to work on interesting projects. Send me an email at ramu.narasinga@gmail.com

My Github — https://github.com/ramu-narasinga

My website — https://ramunarasinga.com

My Youtube channel — https://www.youtube.com/@ramu-narasinga

Learning platform — https://thinkthroo.com

Codebase Architecture — https://app.thinkthroo.com/architecture

Best practices — https://app.thinkthroo.com/best-practices

Production-grade projects — https://app.thinkthroo.com/production-grade-projects

References:

1. https://github.com/vercel/ai-chatbot/blob/main/lib/ai/tools/create-document.ts#L16

2. https://github.com/vercel/ai-chatbot/blob/main/lib/ai/tools/create-document.ts#L2

3. https://www.npmjs.com/package/ai

4. https://github.com/vercel/ai/blob/main/packages/ai/core/tool/tool.ts