Include agent schemas in deployment manifests

sagentic-ai · Feb 26, 2024 · 079f7eb · 079f7eb
1 parent 064f242
commit 079f7eb
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 4 deletions.
diff --git a/bin/bazed.ts b/bin/bazed.ts
@@ -17,6 +17,8 @@ import FormData from "form-data";
 import { SingleBar } from "cli-progress";
 import tar from "tar";
 import moment from "moment";
+import { getToolInterface } from "../src/tool";
+import zodToJsonSchema from "zod-to-json-schema";
 
 dotenv.config();
 
@@ -352,7 +354,22 @@ const scanForAgents = async (path: string): Promise<Record<string, any>> => {
   const constructors = constructorsFromModule(module);
   const agents: Record<string, any> = {};
   for (const constructor of constructors) {
-    agents[constructor.name] = { input: {}, output: {}, state: {} };
+    const toolInterface = getToolInterface(constructor);
+    if (toolInterface) {
+      agents[constructor.name] = {
+        description: toolInterface.description,
+        input: zodToJsonSchema(toolInterface.args),
+        output: zodToJsonSchema(toolInterface.returns),
+        state: {},
+      };
+    } else {
+      agents[constructor.name] = {
+        description: "",
+        input: {},
+        output: {},
+        state: {},
+      };
+    }
   }
   return agents;
 };
@@ -555,14 +572,18 @@ program
           console.log("\tCost per model:");
           for (const model in response.data.session.cost) {
             console.log(
-              `\t\t${model}: $${response.data.session.cost[model as ModelType].toFixed(2)}`,
+              `\t\t${model}: $${response.data.session.cost[
+                model as ModelType
+              ].toFixed(2)}`,
               chalk.gray(`(${response.data.session.cost[model as ModelType]})`)
             );
           }
           console.log("\tTokens per model:");
           for (const model in response.data.session.tokens) {
             console.log(
-              `\t\t${model}: ${response.data.session.tokens[model as ModelType]}`
+              `\t\t${model}: ${
+                response.data.session.tokens[model as ModelType]
+              }`
             );
           }
         }

diff --git a/src/tool.ts b/src/tool.ts
@@ -94,12 +94,21 @@ export class FunctionTool<Args, Returns> implements Tool {
   }
 }
 
+interface ToolInterface {
+  name: string;
+  description: string;
+  args: z.ZodType<any>;
+  returns: z.ZodType<any>;
+}
+
 interface ToolableAgentConstructor extends Constructor<Agent> {
+  toolInterface?: ToolInterface;
   asTool?: Tool;
 }
 
 /** Toolable is a mixin for agent classes that can be used as tools. */
 export interface Toolable {
+  toolInterface?: ToolInterface;
   asTool?: Tool;
 }
 
@@ -122,6 +131,21 @@ export const toTool = (toolLike: ToolLike): Tool => {
   }
 };
 
+/** Get the tool interface of a toolable.
+ * @param toolLike
+ * @returns Tool interface or undefined
+ */
+export const getToolInterface = (
+  toolLike: ToolLike
+): ToolInterface | undefined => {
+  if ("toolInterface" in toolLike) {
+    if (toolLike.toolInterface) {
+      return toolLike.toolInterface;
+    }
+  }
+  return undefined;
+};
+
 /** Decorator for agent classes that can be used as tools.
  * @param name Name of the tool
  * @param args Zod schema for parameters of the tool
@@ -134,7 +158,14 @@ export const isTool = <A, R>(
   returns: z.ZodType<R>
 ) =>
   function (constructor: Constructor<Agent>, _context: any) {
-    (constructor as ToolableAgentConstructor).asTool = new FunctionTool(
+    const c = constructor as ToolableAgentConstructor;
+    c.toolInterface = {
+      name,
+      description,
+      args,
+      returns,
+    };
+    c.asTool = new FunctionTool(
       name,
       description,
       args,