1 files changed, 44 insertions, 23 deletions
diff --git a/js/src/tool/tool.ts b/js/src/tool/tool.ts
index 08ab935b3..748025707 100644
--- a/js/src/tool/tool.ts
+++ b/js/src/tool/tool.ts
@@ -1,28 +1,49 @@
-import { type Tool, tool as AITool } from "ai";
+import { tool, type Tool as AITool } from "ai";
 import { Log } from "../util/log";
 
 const log = Log.create({ service: "tool" });
 
-export function tool<Params, Result>(
-  tool: Tool<Params, Result> & {
-    name: string;
-  },
-) {
-  return {
-    [tool.name]: AITool({
-      ...tool,
-      execute: async (params, opts) => {
-        log.info("invoking", {
-          id: opts.toolCallId,
-          name: tool.name,
-          ...params,
-        });
-        try {
-          return tool.execute!(params, opts);
-        } catch (e: any) {
-          return "An error occurred: " + e.toString();
-        }
-      },
-    }),
-  };
+export namespace Tool {
+  export interface Metadata {
+    properties: Record<string, any>;
+    time: {
+      start: number;
+      end: number;
+    };
+  }
+  export function define<Params, Output>(
+    input: AITool<Params, { metadata?: any; output: Output }> & {
+      name: string;
+    },
+  ) {
+    return {
+      [input.name]: tool({
+        ...input,
+        execute: async (params, opts) => {
+          log.info("invoking", {
+            id: opts.toolCallId,
+            name: input.name,
+            ...params,
+          });
+          try {
+            const start = Date.now();
+            const result = await input.execute!(params, opts);
+            const metadata: Metadata = {
+              properties: result.metadata,
+              time: {
+                start,
+                end: Date.now(),
+              },
+            };
+            return {
+              metadata,
+              output: result.output,
+            };
+          } catch (e: any) {
+            return "An error occurred: " + e.toString();
+          }
+        },
+      }),
+    };
+  }
 }