feat: add type to tool maps

mlc-ai · Oct 22, 2024 · 4e78dbe · 4e78dbe
1 parent a2839e4
commit 4e78dbe
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/tool.ts b/src/tool.ts
@@ -3,8 +3,8 @@ import { retrievers } from "./retriever";
 import { State } from "./state";
 
 export const tool: Record<string, Tool> = {
-  ...retrievers,
-  ...actions,
+  ...Object.fromEntries(Object.entries(retrievers).map(([key, value]) => [key, { ...value, type: "retriever" }])),
+  ...Object.fromEntries(Object.entries(actions).map(([key, value]) => [key, { ...value, type: "action" }])),
 };
 
 export const toolName = Object.keys(tool);
@@ -28,5 +28,6 @@ export interface Tool {
       };
     };
   };
+  type: 'retriever' | 'action';
   implementation: (state: State, parameters: any) => void;
 }