fix: tool serialisation

i-am-bee · Sep 20, 2024 · e0e9095 · e0e9095
1 parent 0c3466b
commit e0e9095
Showing 1 changed file with 23 additions and 0 deletions.
diff --git a/src/tools/imageDescription.ts b/src/tools/imageDescription.ts
@@ -84,6 +84,29 @@ export class ImageDescriptionTool extends Tool<StringToolOutput, ToolOptions, To
     );
   }
 
+  createSnapshot() {
+    return {
+      ...super.createSnapshot(),
+      vllmApiEndpoint: this.vllmApiEndpoint,
+      vllmApiModelId: this.vllmApiModelId,
+      openApiKey: this.openApiKey, // pragma: allowlist secret
+    };
+  }
+
+  loadSnapshot({
+    vllmApiEndpoint,
+    vllmApiModelId,
+    openApiKey,
+    ...snapshot
+  }: ReturnType<typeof this.createSnapshot>) {
+    super.loadSnapshot(snapshot);
+    Object.assign(this, {
+      vllmApiEndpoint,
+      vllmApiModelId,
+      openApiKey,
+    });
+  }
+
   protected async queryVllmAPI(completionPrompt: VllmChatCompletionPrompt) {
     const vllmApiUrl = new URL("/v1/chat/completions", this.vllmApiEndpoint);
     const headers = {