feat: additional tool argument (#1693)

.changeset/ready-lines-rescue.md

@@ -0,0 +1,6 @@
+---
+"@llamaindex/core": patch
+"@llamaindex/doc": patch
+---
+
+Support binding additional argument to function tool

apps/next/src/content/docs/llamaindex/modules/tool/index.mdx

@@ -0,0 +1,60 @@
+---
+title: Tools
+---
+
+A "tool" is a utility that can be called by an agent on behalf of an LLM.
+A tool can be called to perform custom actions, or retrieve extra information based on the LLM-generated input.
+A result from a tool call can be used by subsequent steps in a workflow, or to compute a final answer.
+For example, a "weather tool" could fetch some live weather information from a geographical location.
+
+## Function tool
+
+Function tools are implemented with the `FunctionTool` class.
+A `FunctionTool` is constructed from a function with signature
+```ts
+(input: T, additionalArg?: AdditionalToolArgument) => R
+```
+where
+- `input` is generated by the LLM, `T` is the type defined by the tool `parameters`
+- `additionalArg` is an optional extra argument, see "Binding" below
+- `R` is the return type
+
+### Binding
+
+An additional argument can be bound to a tool, each tool call will be passed
+- the input provided by the LLM
+- the additional argument (extends object)
+
+Note: calling the `bind` method will return a new `FunctionTool` instance, without modifying the tool which `bind` is called on.
+
+Example to pass a `userToken` as additional argument:
+```ts
+// first arg is LLM input, second is bound arg
+const queryKnowledgeBase = async ({ question }, { userToken }) => {
+  const response = await fetch(`https://knowledge-base.com?token=${userToken}&query=${question}`);
+  // ...
+};
+
+// define tool as usual
+const kbTool = FunctionTool.from(queryKnowledgeBase, {
+  name: 'queryKnowledgeBase',
+  description: 'Query knowledge base',
+  parameters: {
+    type: 'object',
+    properties: {
+      question: {
+        type: 'string',
+        description: 'The user question',
+      },
+    },
+    required: ['question'],
+  },
+});
+
+// create an agent
+const additionalArg = { userToken: 'abcd1234' };
+const kbAgent = new LLMAgent({
+  tools: [kbTool.bind(additionalArg)],
+  // llm, systemPrompt etc
+})
+```

packages/core/src/llms/type.ts

@@ -224,8 +224,10 @@ export type ToolMetadata<
 /**
  * Simple Tool interface. Likely to change.
  */
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-export interface BaseTool<Input = any> {
+export interface BaseTool<
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  Input = any,
+> {
   /**
    * This could be undefined if the implementation is not provided,
    *  which might be the case when communicating with a llm.

packages/core/src/tools/function-tool.ts

@@ -4,40 +4,66 @@ import { zodToJsonSchema } from "zod-to-json-schema";
 import type { JSONValue } from "../global";
 import type { BaseTool, ToolMetadata } from "../llms";
 
-export class FunctionTool<T, R extends JSONValue | Promise<JSONValue>>
-  implements BaseTool<T>
+export class FunctionTool<
+  T,
+  R extends JSONValue | Promise<JSONValue>,
+  AdditionalToolArgument extends object = object,
+> implements BaseTool<T>
 {
-  #fn: (input: T) => R;
+  #fn: (input: T, additionalArg?: AdditionalToolArgument) => R;
+  #additionalArg: AdditionalToolArgument | undefined;
   readonly #metadata: ToolMetadata<JSONSchemaType<T>>;
   readonly #zodType: z.ZodType<T> | null = null;
   constructor(
-    fn: (input: T) => R,
+    fn: (input: T, additionalArg?: AdditionalToolArgument) => R,
     metadata: ToolMetadata<JSONSchemaType<T>>,
     zodType?: z.ZodType<T>,
+    additionalArg?: AdditionalToolArgument,
   ) {
     this.#fn = fn;
     this.#metadata = metadata;
     if (zodType) {
       this.#zodType = zodType;
     }
+    this.#additionalArg = additionalArg;
   }
 
-  static from<T>(
-    fn: (input: T) => JSONValue | Promise<JSONValue>,
+  static from<T, AdditionalToolArgument extends object = object>(
+    fn: (
+      input: T,
+      additionalArg?: AdditionalToolArgument,
+    ) => JSONValue | Promise<JSONValue>,
     schema: ToolMetadata<JSONSchemaType<T>>,
-  ): FunctionTool<T, JSONValue | Promise<JSONValue>>;
-  static from<R extends z.ZodType>(
-    fn: (input: z.infer<R>) => JSONValue | Promise<JSONValue>,
+  ): FunctionTool<T, JSONValue | Promise<JSONValue>, AdditionalToolArgument>;
+  static from<
+    R extends z.ZodType,
+    AdditionalToolArgument extends object = object,
+  >(
+    fn: (
+      input: z.infer<R>,
+      additionalArg?: AdditionalToolArgument,
+    ) => JSONValue | Promise<JSONValue>,
     schema: Omit<ToolMetadata, "parameters"> & {
       parameters: R;
     },
-  ): FunctionTool<z.infer<R>, JSONValue | Promise<JSONValue>>;
-  static from<T, R extends z.ZodType<T>>(
-    fn: (input: T) => JSONValue | Promise<JSONValue>,
+  ): FunctionTool<
+    z.infer<R>,
+    JSONValue | Promise<JSONValue>,
+    AdditionalToolArgument
+  >;
+  static from<
+    T,
+    R extends z.ZodType<T>,
+    AdditionalToolArgument extends object = object,
+  >(
+    fn: (
+      input: T,
+      additionalArg?: AdditionalToolArgument,
+    ) => JSONValue | Promise<JSONValue>,
     schema: Omit<ToolMetadata, "parameters"> & {
       parameters: R;
     },
-  ): FunctionTool<T, JSONValue>;
+  ): FunctionTool<T, JSONValue, AdditionalToolArgument>;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   static from(fn: any, schema: any): any {
     if (schema.parameters instanceof z.ZodSchema) {
@@ -58,6 +84,15 @@ export class FunctionTool<T, R extends JSONValue | Promise<JSONValue>>
     return this.#metadata as BaseTool<T>["metadata"];
   }
 
+  bind = (additionalArg: AdditionalToolArgument) => {
+    return new FunctionTool(
+      this.#fn,
+      this.#metadata,
+      this.#zodType ?? undefined,
+      additionalArg,
+    );
+  };
+
   call = (input: T) => {
     if (this.#metadata.requireContext) {
       const inputWithContext = input as Record<string, unknown>;
@@ -72,15 +107,18 @@ export class FunctionTool<T, R extends JSONValue | Promise<JSONValue>>
       if (result.success) {
         if (this.#metadata.requireContext) {
           const { context } = input as Record<string, unknown>;
-          return this.#fn.call(null, { context, ...result.data });
+          return this.#fn.call(
+            null,
+            { context, ...result.data },
+            this.#additionalArg,
+          );
         } else {
-          return this.#fn.call(null, result.data);
+          return this.#fn.call(null, result.data, this.#additionalArg);
         }
       } else {
         console.warn(result.error.errors);
       }
     }
-
-    return this.#fn.call(null, input);
+    return this.#fn.call(null, input, this.#additionalArg);
   };
 }

packages/core/tests/tools.test.ts

@@ -1,5 +1,5 @@
 import { FunctionTool } from "@llamaindex/core/tools";
-import { describe, test } from "vitest";
+import { describe, expect, test, vi } from "vitest";
 import { z } from "zod";
 
 describe("FunctionTool", () => {
@@ -32,4 +32,38 @@ describe("FunctionTool", () => {
       parameters: inputSchema,
     });
   });
+
+  test("bind additional argument", () => {
+    type AdditionalHelloArgument = {
+      question?: string;
+    };
+
+    const hello = vi
+      .fn()
+      .mockImplementation((name: string, arg?: AdditionalHelloArgument) => {
+        return `Hello ${name}. ${arg?.question ?? ""}`;
+      });
+
+    const helloTool = FunctionTool.from<string, AdditionalHelloArgument>(
+      hello,
+      {
+        name: "hello",
+        description: "Says hello",
+      },
+    );
+
+    helloTool.call("Alice");
+    expect(hello).to.toHaveBeenCalledOnce();
+    expect(hello).to.toHaveBeenCalledWith("Alice", undefined);
+
+    hello.mockReset();
+
+    const additionalArg = {
+      question: "How is it going?",
+    };
+    const helloBoundTool = helloTool.bind(additionalArg);
+    helloBoundTool.call("Bob");
+    expect(hello).to.toHaveBeenCalledOnce();
+    expect(hello).to.toHaveBeenCalledWith("Bob", additionalArg);
+  });
 });