Allow tools to be initialized with JSON schema

langchain-core/src/tools/index.ts

@@ -20,6 +20,7 @@ import { ZodObjectAny } from "../types/zod.js";
 import { MessageContent } from "../messages/base.js";
 import { AsyncLocalStorageProviderSingleton } from "../singletons/index.js";
 import { _isToolCall, ToolInputParsingException } from "./utils.js";
+import { isZodSchema } from "../utils/types/is_zod_schema.js";
 
 export { ToolInputParsingException };
 
@@ -319,16 +320,18 @@ export interface DynamicToolInput extends BaseDynamicToolInput {
  * Interface for the input parameters of the DynamicStructuredTool class.
  */
 export interface DynamicStructuredToolInput<
-  T extends ZodObjectAny = ZodObjectAny
+  T extends ZodObjectAny | Record<string, any> = ZodObjectAny
 > extends BaseDynamicToolInput {
   func: (
     input: BaseDynamicToolInput["responseFormat"] extends "content_and_artifact"
       ? ToolCall
-      : z.infer<T>,
+      : T extends ZodObjectAny
+      ? z.infer<T>
+      : T,
     runManager?: CallbackManagerForToolRun,
     config?: RunnableConfig
   ) => Promise<ToolReturnType>;
-  schema: T;
+  schema: T extends ZodObjectAny ? T : T;
 }
 
 /**
@@ -384,8 +387,8 @@ export class DynamicTool extends Tool {
  * provided function when the tool is called.
  */
 export class DynamicStructuredTool<
-  T extends ZodObjectAny = ZodObjectAny
-> extends StructuredTool<T> {
+  T extends ZodObjectAny | Record<string, any> = ZodObjectAny
+> extends StructuredTool<T extends ZodObjectAny ? T : ZodObjectAny> {
   static lc_name() {
     return "DynamicStructuredTool";
   }
@@ -396,22 +399,24 @@ export class DynamicStructuredTool<
 
   func: DynamicStructuredToolInput<T>["func"];
 
-  schema: T;
+  schema: T extends ZodObjectAny ? T : ZodObjectAny;
 
   constructor(fields: DynamicStructuredToolInput<T>) {
     super(fields);
     this.name = fields.name;
     this.description = fields.description;
     this.func = fields.func;
     this.returnDirect = fields.returnDirect ?? this.returnDirect;
-    this.schema = fields.schema;
+    this.schema = (
+      isZodSchema(fields.schema) ? fields.schema : z.object({})
+    ) as T extends ZodObjectAny ? T : ZodObjectAny;
   }
 
   /**
    * @deprecated Use .invoke() instead. Will be removed in 0.3.0.
    */
   async call(
-    arg: z.output<T> | ToolCall,
+    arg: (T extends ZodObjectAny ? z.output<T> : T) | ToolCall,
     configArg?: RunnableConfig | Callbacks,
     /** @deprecated */
     tags?: string[]
@@ -424,11 +429,11 @@ export class DynamicStructuredTool<
   }
 
   protected _call(
-    arg: z.output<T> | ToolCall,
+    arg: (T extends ZodObjectAny ? z.output<T> : T) | ToolCall,
     runManager?: CallbackManagerForToolRun,
     parentConfig?: RunnableConfig
   ): Promise<ToolReturnType> {
-    return this.func(arg, runManager, parentConfig);
+    return this.func(arg as any, runManager, parentConfig);
   }
 }
 
@@ -447,10 +452,13 @@ export abstract class BaseToolkit {
 
 /**
  * Parameters for the tool function.
- * @template {ZodObjectAny | z.ZodString = ZodObjectAny} RunInput The input schema for the tool. Either any Zod object, or a Zod string.
+ * @template {ZodObjectAny | z.ZodString | Record<string, any> = ZodObjectAny} RunInput The input schema for the tool. Either any Zod object, a Zod string, or JSON schema.
  */
 interface ToolWrapperParams<
-  RunInput extends ZodObjectAny | z.ZodString = ZodObjectAny
+  RunInput extends
+    | ZodObjectAny
+    | z.ZodString
+    | Record<string, any> = ZodObjectAny
 > extends ToolParams {
   /**
    * The name of the tool. If using with an LLM, this
@@ -494,18 +502,25 @@ interface ToolWrapperParams<
  *
  * @returns {DynamicStructuredTool<T>} A new StructuredTool instance.
  */
-export function tool<T extends z.ZodString = z.ZodString>(
+export function tool<T extends z.ZodString>(
   func: RunnableFunc<z.output<T>, ToolReturnType>,
   fields: ToolWrapperParams<T>
 ): DynamicTool;
 
-export function tool<T extends ZodObjectAny = ZodObjectAny>(
+export function tool<T extends ZodObjectAny>(
   func: RunnableFunc<z.output<T>, ToolReturnType>,
   fields: ToolWrapperParams<T>
 ): DynamicStructuredTool<T>;
 
-export function tool<T extends ZodObjectAny | z.ZodString = ZodObjectAny>(
-  func: RunnableFunc<z.output<T>, ToolReturnType>,
+export function tool<T extends Record<string, any>>(
+  func: RunnableFunc<T, ToolReturnType>,
+  fields: ToolWrapperParams<T>
+): DynamicStructuredTool<T>;
+
+export function tool<
+  T extends ZodObjectAny | z.ZodString | Record<string, any> = ZodObjectAny
+>(
+  func: RunnableFunc<T extends ZodObjectAny ? z.output<T> : T, ToolReturnType>,
   fields: ToolWrapperParams<T>
 ):
   | DynamicStructuredTool<T extends ZodObjectAny ? T : ZodObjectAny>
@@ -518,7 +533,7 @@ export function tool<T extends ZodObjectAny | z.ZodString = ZodObjectAny>(
         fields.description ??
         fields.schema?.description ??
         `${fields.name} tool`,
-      func,
+      func: func as any,
     });
   }
 
@@ -528,7 +543,7 @@ export function tool<T extends ZodObjectAny | z.ZodString = ZodObjectAny>(
   return new DynamicStructuredTool<T extends ZodObjectAny ? T : ZodObjectAny>({
     ...fields,
     description,
-    schema: fields.schema as T extends ZodObjectAny ? T : ZodObjectAny,
+    schema: fields.schema as any,
     // TODO: Consider moving into DynamicStructuredTool constructor
     func: async (input, runManager, config) => {
       return new Promise((resolve, reject) => {
@@ -539,7 +554,7 @@ export function tool<T extends ZodObjectAny | z.ZodString = ZodObjectAny>(
           childConfig,
           async () => {
             try {
-              resolve(func(input, childConfig));
+              resolve(func(input as any, childConfig));
             } catch (e) {
               reject(e);
             }

langchain-core/src/tools/tests/tools.test.ts

@@ -1,6 +1,6 @@
 import { test, expect } from "@jest/globals";
 import { z } from "zod";
-import { tool } from "../index.js";
+import { DynamicStructuredTool, tool } from "../index.js";
 import { ToolMessage } from "../../messages/tool.js";
 
 test("Tool should error if responseFormat is content_and_artifact but the function doesn't return a tuple", async () => {
@@ -115,3 +115,100 @@ test("Tool can accept single string input", async () => {
   const result = await stringTool.invoke("b");
   expect(result).toBe("ba");
 });
+
+test("Tool declared with JSON schema", async () => {
+  const weatherSchema = {
+    type: "object",
+    properties: {
+      location: {
+        type: "string",
+        description: "A place",
+      },
+    },
+    required: ["location"],
+  };
+  const weatherTool = tool(
+    (_) => {
+      return "Sunny";
+    },
+    {
+      name: "weather",
+      schema: weatherSchema,
+    }
+  );
+
+  const weatherTool2 = new DynamicStructuredTool({
+    name: "weather",
+    description: "get the weather",
+    func: async (_) => {
+      return "Sunny";
+    },
+    schema: weatherSchema,
+  });
+  // No validation on JSON schema tools
+  await weatherTool.invoke({
+    somethingSilly: true,
+  });
+  await weatherTool2.invoke({
+    somethingSilly: true,
+  });
+});
+
+test("Tool input typing is enforced", async () => {
+  const weatherSchema = z.object({
+    location: z.string(),
+  });
+
+  const weatherTool = tool(
+    (_) => {
+      return "Sunny";
+    },
+    {
+      name: "weather",
+      schema: weatherSchema,
+    }
+  );
+
+  const weatherTool2 = new DynamicStructuredTool({
+    name: "weather",
+    description: "get the weather",
+    func: async (_) => {
+      return "Sunny";
+    },
+    schema: weatherSchema,
+  });
+
+  const weatherTool3 = tool(
+    async (_) => {
+      return "Sunny";
+    },
+    {
+      name: "weather",
+      description: "get the weather",
+      schema: z.string(),
+    }
+  );
+
+  await expect(async () => {
+    await weatherTool.invoke({
+      // @ts-expect-error Invalid argument
+      badval: "someval",
+    });
+  }).rejects.toThrow();
+  const res = await weatherTool.invoke({
+    location: "somewhere",
+  });
+  expect(res).toEqual("Sunny");
+  await expect(async () => {
+    await weatherTool2.invoke({
+      // @ts-expect-error Invalid argument
+      badval: "someval",
+    });
+  }).rejects.toThrow();
+  const res2 = await weatherTool2.invoke({
+    location: "someval",
+  });
+  expect(res2).toEqual("Sunny");
+  const res3 = await weatherTool3.invoke("blah");
+  expect(res3).toEqual("Sunny");
+});