feat(llm): make drivers serializable

i-am-bee · Sep 23, 2024 · 05d6dbb · 05d6dbb
1 parent a91b8db
commit 05d6dbb
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 2 deletions.
diff --git a/src/llms/drivers/base.ts b/src/llms/drivers/base.ts
@@ -27,13 +27,16 @@ import { Retryable } from "@/internals/helpers/retryable.js";
 import { PromptTemplate } from "@/template.js";
 import { SchemaObject } from "ajv";
 import { z } from "zod";
+import { Serializable } from "@/internals/serializable.js";
 
 export interface GenerateSchemaInput<T> {
   maxRetries?: number;
   options?: T;
 }
 
-export abstract class BaseDriver<TGenerateOptions extends GenerateOptions = GenerateOptions> {
+export abstract class BaseDriver<
+  TGenerateOptions extends GenerateOptions = GenerateOptions,
+> extends Serializable<any> {
   protected abstract template: PromptTemplate.infer<{ schema: string }>;
   protected errorTemplate = new PromptTemplate({
     schema: z.object({
@@ -45,7 +48,9 @@ export abstract class BaseDriver<TGenerateOptions extends GenerateOptions = Gene
 Validation Errors: "{{errors}}"`,
   });
 
-  constructor(protected llm: ChatLLM<ChatLLMOutput, TGenerateOptions>) {}
+  constructor(protected readonly llm: ChatLLM<ChatLLMOutput, TGenerateOptions>) {
+    super();
+  }
 
   protected abstract parseResponse(textResponse: string): unknown;
   protected abstract schemaToString(schema: SchemaObject): Promise<string> | string;
@@ -123,4 +128,15 @@ Validation Errors: "{{errors}}"`,
       },
     }).get();
   }
+
+  createSnapshot() {
+    return {
+      template: this.template,
+      errorTemplate: this.errorTemplate,
+    };
+  }
+
+  loadSnapshot(snapshot: ReturnType<typeof this.createSnapshot>) {
+    Object.assign(this, snapshot);
+  }
 }
diff --git a/src/llms/drivers/json.ts b/src/llms/drivers/json.ts
@@ -38,6 +38,10 @@ IMPORTANT: Every message must be a parsable JSON string without additional outpu
 `,
   });
 
+  static {
+    this.register();
+  }
+
   protected parseResponse(textResponse: string): unknown {
     return parseBrokenJson(textResponse);
   }

diff --git a/src/llms/drivers/typescript.ts b/src/llms/drivers/typescript.ts
@@ -39,6 +39,10 @@ IMPORTANT: Every message must be a parsable JSON string without additional outpu
 `,
   });
 
+  static {
+    this.register();
+  }
+
   protected parseResponse(textResponse: string): unknown {
     return parseBrokenJson(textResponse);
   }

diff --git a/src/llms/drivers/yaml.ts b/src/llms/drivers/yaml.ts
@@ -38,6 +38,10 @@ IMPORTANT: Every message must be a parsable YAML string without additional outpu
 `,
   });
 
+  static {
+    this.register();
+  }
+
   protected parseResponse(textResponse: string): unknown {
     return yaml.load(textResponse);
   }