ToolBuilder

Author: BLamy

src/Chat.ts

@@ -1,8 +1,10 @@
 import OpenAI from "openai";
 import { PromptBuilder } from "./PromptBuilder";
 import { ExtractArgs, ExtractChatArgs, ReplaceChatArgs } from "./types";
+import { ToolBuilder } from "./ToolBuilder";
 
 export class Chat<
+  ToolNames extends string,
   TMessages extends
     | []
     | [
@@ -14,8 +16,13 @@ export class Chat<
   constructor(
     public messages: TMessages,
     public args: TSuppliedInputArgs,
+    public tools = {} as Record<ToolNames, ToolBuilder>,
+    public mustUseTool: boolean = false
   ) {}
 
+  toJSONSchema() {
+  }
+
   toArray() {
     return this.messages.map((m) => ({
       role: m.role,

src/ToolBuilder.ts

@@ -0,0 +1,51 @@
+interface Tool<I = unknown, O = unknown> {
+    name: string;
+    type: "query" | "mutation"
+    build: (input: I) => O;
+}
+
+export class ToolBuilder<TType extends "query" | "mutation" = "query", I = unknown, O = unknown> {
+    private name: string;
+    private implementation?: (input: I) => O;
+    private type: TType;
+  
+    constructor(name: string, type: TType = "query" as TType) {
+      this.name = name;
+      this.type = type;
+    }
+  
+    addInputValidation<T = I>(): ToolBuilder<TType, T, O> {
+      // Implementation here
+      return this as unknown as ToolBuilder<TType, T, O>;
+    }
+  
+    addOutputValidation<T = O>(): ToolBuilder<TType, I, T> {
+      // Implementation here
+      return this as unknown as ToolBuilder<TType, I, T>;
+    }
+  
+    query(queryFunction: (input: I) => O): ToolBuilder<"query", I, O> {
+
+      return {
+        ...this,
+        implementation: queryFunction,
+        type: "query"
+        };    
+    }
+  
+    mutation(mutationFunction: (input: I) => O): ToolBuilder<"mutation", I, O> {
+      return {
+        ...this,
+        implementation: mutationFunction,
+        type: "mutation"
+        };
+    }
+  
+    build(): Tool<I, O> {
+      return {
+        name: this.name,
+        build: this.implementation!,
+        type: this.type
+      };
+    }
+  }

src/__tests__/Chat.test.ts

@@ -2,10 +2,11 @@ import { strict as assert } from "node:assert";
 import { Chat } from "../Chat";
 import { system, user, assistant } from "../ChatHelpers";
 import { Equal, Expect } from "./types.test";
+import { ToolBuilder } from "../ToolBuilder";
 
 describe("Chat", () => {
   it("should allow empty array", () => {
-    const chat = new Chat([], {}).toArray();
+    const chat = new Chat([], {}, {}).toArray();
     type test = Expect<Equal<typeof chat, []>>;
     assert.deepEqual(chat, []);
   });
@@ -54,16 +55,97 @@ describe("Chat", () => {
   });
 
   it("should allow chat of all diffent types with no args", () => {
-    const chat = new Chat(
-      [user(`Tell me a joke`), assistant(`joke?`), system(`joke?`)],
-      {},
-    ).toArray();
     const usrMsg = user("Tell me a joke");
     const astMsg = assistant("joke?");
     const sysMsg = system("joke?");
+
+    const chat = new Chat([usrMsg, astMsg, sysMsg], {}).toArray();
     type test = Expect<
       Equal<typeof chat, [typeof usrMsg, typeof astMsg, typeof sysMsg]>
     >;
     assert.deepEqual(chat, [usrMsg, astMsg, sysMsg]);
   });
+
+  it("should allow me to pass in tools", () => {
+    const usrMsg = user("Tell me a joke");
+    const astMsg = assistant("joke?");
+    const sysMsg = system("joke?");
+    const tools = {
+      google: new ToolBuilder("google")
+        .addInputValidation<{ query: string }>()
+        .addOutputValidation<{ results: string[] }>()
+        .query(({ query }) => {
+          return {
+            results: ["foo", "bar"],
+          };
+        }),
+      wikipedia: new ToolBuilder("wikipedia")
+        .addInputValidation<{ page: string }>()
+        .addOutputValidation<{ results: string[] }>()
+        .query(({ page }) => {
+          return {
+            results: ["foo", "bar"],
+          };
+        }),
+      sendEmail: new ToolBuilder("sendEmail")
+        .addInputValidation<{ to: string; subject: string; body: string }>()
+        .addOutputValidation<{ success: boolean }>()
+        .mutation(({ to, subject, body }) => {
+          return {
+            success: true,
+          };
+        }),
+    } as const;
+
+    const chat = new Chat([usrMsg, astMsg, sysMsg], {}, tools);
+
+    type tests = [
+      Expect<
+        Equal<
+          typeof chat,
+          Chat<
+            keyof typeof tools,
+            [typeof usrMsg, typeof astMsg, typeof sysMsg],
+            {}
+          >
+        >
+      >,
+      Expect<
+        Equal<
+          typeof tools,
+          {
+            readonly google: ToolBuilder<
+              "query",
+              {
+                query: string;
+              },
+              {
+                results: string[];
+              }
+            >;
+            readonly wikipedia: ToolBuilder<
+              "query",
+              {
+                page: string;
+              },
+              {
+                results: string[];
+              }
+            >;
+            readonly sendEmail: ToolBuilder<
+              "mutation",
+              {
+                to: string;
+                subject: string;
+                body: string;
+              },
+              {
+                success: boolean;
+              }
+            >;
+          }
+        >
+      >
+    ];
+  });
 });