rework custom tools

2 files changed, 11 insertions, 12 deletions

diff --git a/packages/plugin/src/example.ts b/packages/plugin/src/example.ts
index fd6a404d9..1e4557a68 100644
--- a/packages/plugin/src/example.ts
+++ b/packages/plugin/src/example.ts
@@ -5,15 +5,15 @@ export const ExamplePlugin: Plugin = async (ctx) => {
   return {
     permission: {},
     tool: {
-      mytool: tool((zod) => ({
+      mytool: tool({
         description: "This is a custom tool tool",
         args: {
-          foo: zod.string(),
+          foo: tool.schema.string().describe("foo"),
         },
-        async execute(args, ctx) {
+        async execute(args) {
           return `Hello ${args.foo}!`
         },
-      })),
+      }),
     },
     async "chat.params"(_input, output) {
       output.topP = 1
diff --git a/packages/plugin/src/tool.ts b/packages/plugin/src/tool.ts
index 7c1d3d7c5..2998a1e72 100644
--- a/packages/plugin/src/tool.ts
+++ b/packages/plugin/src/tool.ts
@@ -7,14 +7,13 @@ export type ToolContext = {
   abort: AbortSignal
 }
 
-export function tool<Args extends z.ZodRawShape>(
-  input: (zod: typeof z) => {
-    description: string
-    args: Args
-    execute: (args: z.infer<z.ZodObject<Args>>, ctx: ToolContext) => Promise<string>
-  },
-) {
-  return input(z)
+export function tool<Args extends z.ZodRawShape>(input: {
+  description: string
+  args: Args
+  execute(args: z.infer<z.ZodObject<Args>>, context: ToolContext): Promise<string>
+}) {
+  return input
 }
+tool.schema = z
 
 export type ToolDefinition = ReturnType<typeof tool>