From e0e9095073040cd4545e27671e1322ea670964ba Mon Sep 17 00:00:00 2001
From: Graham White <gwhite@uk.ibm.com>
Date: Fri, 20 Sep 2024 11:57:49 +0100
Subject: [PATCH] fix: tool serialisation

---
 src/tools/imageDescription.ts | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/src/tools/imageDescription.ts b/src/tools/imageDescription.ts
index adaf328..e9307f3 100644
--- 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 = {