[FIX] Fixed types

Author: anuran-roy

src/utils/chat.ts

@@ -4,25 +4,25 @@ import type { z } from "zod";
 import { zodToJsonSchema } from 'zod-to-json-schema';
 import { parseUntilJson } from "./parseUntilJson";
 
-export class ChatPromptTemplate {
+export class ChatPromptTemplate<T> {
   protected template = "";
   protected llm: OpenAI | null = null;
-  protected variables: string[] | null = null;
+  protected variables: (keyof T)[] | null = null;
   protected invokeFn: (...params: any) => string = () => { return "" }
 
   constructor(params: {
     template: string;
-    inputVariables: string[];
+    inputVariables: (keyof T)[];
     templateFormat?: "mustache";
   }) {
     this.template = params.template;
     this.variables = params.inputVariables;
   }
 
-  public static fromTemplate(template: string): ChatPromptTemplate {
+  public static fromTemplate<T extends z.infer<z.ZodObject>>(template: string): ChatPromptTemplate<T> {
     // Extract mustache-style variables from the template
     const variableRegex = /{{\s*([a-zA-Z0-9_]+)\s*}}/g;
-    const variables: string[] = [];
+    const variables: (keyof T)[] = [];
     let match: RegExpExecArray | null = variableRegex.exec(template);
     while (match !== null) {
       variables.push(match[1] ?? "");
@@ -51,20 +51,20 @@ export class ChatPromptTemplate {
     }
   }
 
-  public format(params: Record<string, any>) {
+  public format<T extends z.infer<z.ZodObject>>(params: T) {
     const paramsInVariables = Object.fromEntries(
       Object.keys(params)
-        .filter((key) => (this.variables ?? []).includes(key))
+        .filter((key) => (this.variables ?? []).includes(key as never))
         .map((key) => [key, params[key]]),
     );
     const finalTemplate = Mustache.render(this.template, paramsInVariables);
     return finalTemplate;
   }
 
-  public async invoke(model: string, params: Record<string, any>) {
+  public async invoke<T extends z.infer<z.ZodObject>>(model: string, params: T) {
     let finalTemplate = this.format(params);
 
-    if (!!this.invokeFn) {
+    if (this.invokeFn) {
       finalTemplate = this.invokeFn(finalTemplate);
     }
 
@@ -108,4 +108,4 @@ export function getValidatedOutput<T extends z.ZodTypeAny>(schema: T, data: stri
   const parsedData = parseUntilJson(data);
 
   return schema.safeParse(parsedData);
-}
+}