wip, feat: typed tool

Author: himself65

src/core/athena.ts

@@ -5,19 +5,64 @@ import logger from "../utils/logger.js";
 
 export type Dict<T> = { [key: string]: T };
 
-export interface IAthenaArgument {
-  type: "string" | "number" | "boolean" | "object" | "array";
+type IAthenaArgumentPrimitive = {
+  type: "string" | "number" | "boolean";
   desc: string;
   required: boolean;
-  of?: Dict<IAthenaArgument> | IAthenaArgument;
-}
-
-export interface IAthenaTool {
+};
+
+export type IAthenaArgument =
+  | IAthenaArgumentPrimitive
+  | {
+      type: "object" | "array";
+      desc: string;
+      required: boolean;
+      of: Dict<IAthenaArgument> | IAthenaArgument;
+    };
+type IAthenaArgumentInstance<T extends IAthenaArgument> =
+  T extends IAthenaArgumentPrimitive
+    ? T["type"] extends "string"
+      ? T["required"] extends true
+        ? string
+        : string | undefined
+      : T["type"] extends "number"
+        ? T["required"] extends true
+          ? number
+          : number | undefined
+        : T["type"] extends "boolean"
+          ? T["required"] extends true
+            ? boolean
+            : boolean | undefined
+          : never
+    : T extends { of: Dict<IAthenaArgument> }
+      ? T["required"] extends true
+        ? { [K in keyof T["of"]]: IAthenaArgumentInstance<T["of"][K]> }
+        :
+            | { [K in keyof T["of"]]: IAthenaArgumentInstance<T["of"][K]> }
+            | undefined
+      : T extends { of: IAthenaArgument }
+        ? T["required"] extends true
+          ? IAthenaArgumentInstance<T["of"]>[]
+          : IAthenaArgumentInstance<T["of"]>[] | undefined
+        : never;
+
+export interface IAthenaTool<
+  Args extends Dict<IAthenaArgument> = Dict<IAthenaArgument>,
+  RetArgs extends Dict<IAthenaArgument> = Dict<IAthenaArgument>,
+> {
   name: string;
   desc: string;
-  args: Dict<IAthenaArgument>;
-  retvals: Dict<IAthenaArgument>;
-  fn: (args: Dict<any>) => Promise<Dict<any>>;
+  args: Args;
+  retvals: RetArgs;
+  fn: (args: {
+    [K in keyof Args]: Args[K] extends IAthenaArgument
+      ? IAthenaArgumentInstance<Args[K]>
+      : never;
+  }) => Promise<{
+    [K in keyof RetArgs]: RetArgs[K] extends IAthenaArgument
+      ? IAthenaArgumentInstance<RetArgs[K]>
+      : never;
+  }>;
   explain_args?: (args: Dict<any>) => IAthenaExplanation;
   explain_retvals?: (args: Dict<any>, retvals: Dict<any>) => IAthenaExplanation;
 }
@@ -38,7 +83,7 @@ export class Athena extends EventEmitter {
   config: Dict<any>;
   states: Dict<Dict<any>>;
   plugins: Dict<PluginBase>;
-  tools: Dict<IAthenaTool>;
+  tools: Dict<IAthenaTool<any, any>>;
   events: Dict<IAthenaEvent>;
 
   constructor(config: Dict<any>, states: Dict<Dict<any>>) {
@@ -103,7 +148,26 @@ export class Athena extends EventEmitter {
     logger.warn(`Plugin ${name} is unloaded`);
   }
 
-  registerTool(tool: IAthenaTool) {
+  registerTool<
+    Args extends Dict<IAthenaArgument>,
+    RetArgs extends Dict<IAthenaArgument>,
+    Tool extends IAthenaTool<Args, RetArgs>,
+  >(
+    config: {
+      name: string;
+      desc: string;
+      args: Args;
+      retvals: RetArgs;
+    },
+    toolImpl: {
+      fn: Tool["fn"];
+      explain_args?: Tool["explain_args"];
+    },
+  ) {
+    const tool = {
+      ...config,
+      ...toolImpl,
+    };
     if (tool.name in this.tools) {
       throw new Error(`Tool ${tool.name} already registered`);
     }

src/plugins/cerebrum/init.ts

@@ -47,36 +47,40 @@ export default class Cerebrum extends PluginBase {
     this.boundAthenaPrivateEventHandler =
       this.athenaPrivateEventHandler.bind(this);
     if (this.config.image_supported) {
-      athena.registerTool({
-        name: "image/check-out",
-        desc: "Check out an image. Whenever you want to see an image, or the user asks you to see an image, use this tool.",
-        args: {
-          image: {
-            type: "string",
-            desc: "The URL or local path of the image to check out.",
-            required: true,
+      athena.registerTool(
+        {
+          name: "image/check-out",
+          desc: "Check out an image. Whenever you want to see an image, or the user asks you to see an image, use this tool.",
+          args: {
+            image: {
+              type: "string",
+              desc: "The URL or local path of the image to check out.",
+              required: true,
+            },
           },
-        },
-        retvals: {
-          result: {
-            type: "string",
-            desc: "The result of checking out the image.",
-            required: true,
+          retvals: {
+            result: {
+              type: "string",
+              desc: "The result of checking out the image.",
+              required: true,
+            },
           },
         },
-        fn: async (args: Dict<any>) => {
-          let image = args.image;
-          if (!image.startsWith("http")) {
-            image = await image2uri(image);
-          }
-          this.imageUrls.push(image);
-          return { result: "success" };
+        {
+          fn: async (args) => {
+            let image = args.image;
+            if (!image.startsWith("http")) {
+              image = await image2uri(image);
+            }
+            this.imageUrls.push(image);
+            return { result: "success" };
+          },
+          explain_args: (args: Dict<any>) => ({
+            summary: "Checking out the image...",
+            details: args.image,
+          }),
         },
-        explain_args: (args: Dict<any>) => ({
-          summary: "Checking out the image...",
-          details: args.image,
-        }),
-      });
+      );
     }
     athena.on("event", this.boundAthenaEventHandler);
     athena.on("private-event", this.boundAthenaPrivateEventHandler);