add tool tests

1 files changed, 40 insertions, 36 deletions

diff --git a/js/src/tool/tool.ts b/js/src/tool/tool.ts
index 47a3918d4..54e4da401 100644
--- a/js/src/tool/tool.ts
+++ b/js/src/tool/tool.ts
@@ -4,49 +4,53 @@ import { Log } from "../util/log";
 const log = Log.create({ service: "tool" });
 
 export namespace Tool {
-  export interface Metadata {
-    properties: Record<string, any>;
+  export interface Metadata<
+    Properties extends Record<string, any> = Record<string, any>,
+  > {
+    properties: Properties;
     time: {
       start: number;
       end: number;
     };
   }
-  export function define<Params, Output>(
-    input: AITool<Params, { metadata?: any; output: Output }> & {
-      name: string;
+  export function define<
+    Params,
+    Output extends { metadata?: any; output: any },
+    Name extends string,
+  >(
+    input: AITool<Params, Output> & {
+      name: Name;
     },
   ) {
-    return {
-      [input.name]: tool({
-        ...input,
-        execute: async (params, opts) => {
-          log.info("invoking", {
-            id: opts.toolCallId,
-            name: input.name,
-            ...params,
+    return 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<Output["metadata"]> = {
+            ...result.metadata,
+            time: {
+              start,
+              end: Date.now(),
+            },
+          };
+          return {
+            metadata,
+            output: result.output,
+          };
+        } catch (e: any) {
+          log.error("error", {
+            msg: e.toString(),
           });
-          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) {
-            log.error("error", {
-              msg: e.toString(),
-            });
-            return "An error occurred: " + e.toString();
-          }
-        },
-      }),
-    };
+          return "An error occurred: " + e.toString();
+        }
+      },
+    });
   }
 }