refactor: image snapshot handling

concord-consortium · Dec 18, 2024 · 392b223 · 392b223
1 parent d0e93ae
commit 392b223
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 14 deletions.
diff --git a/src/models/assistant-model.ts b/src/models/assistant-model.ts
@@ -235,14 +235,17 @@ export const AssistantModel = types
                 const request = { action, resource, values };
                 self.transcriptStore.addMessage(DEBUG_SPEAKER, { description: "Request sent to CODAP", content: formatJsonMessage(request) });
                 let res = yield codapInterface.sendRequest(request);
-                // note: we will implement a new endpoint in the CODAP api to get the exportDataUri value,
-                // so that it can be retrieved separately from a general get component request
-                if (res.values.exportDataUri) {
+                // Prepare for uploading of image file after run if the request is to get dataDisplay
+                const isImageSnapshotRequest = action === "get" && resource.match(/^dataDisplay/);
+                if (isImageSnapshotRequest) {
                   self.uploadFileAfterRun = true;
                   self.dataUri = res.values.exportDataUri;
-                  res = { ...res, values: { ...res.values, exportDataUri: undefined } };
                 }
                 self.transcriptStore.addMessage(DEBUG_SPEAKER, { description: "Response from CODAP", content: formatJsonMessage(res) });
+                // remove any exportDataUri value that exists since it can be large and we don't need to send it to the assistant
+                res = isImageSnapshotRequest
+                  ? { ...res, values: { ...res.values, exportDataUri: undefined } }
+                  : res;
                 return { tool_call_id: toolCall.id, output: JSON.stringify(res) };
               } else {
                 return { tool_call_id: toolCall.id, output: "Tool call not recognized." };

diff --git a/src/utils/openai-utils.ts b/src/utils/openai-utils.ts
@@ -12,18 +12,23 @@ export const newOpenAI = () => {
 };
 
 export async function convertBase64ToImage(base64Data: string, filename = "image.png") {
-  const base64 = base64Data.split(",")[1];
+  try {
+    const mimeType = base64Data.match(/data:(.*?);base64/)?.[1] || "image/png";
+    const base64 = base64Data.split(",")[1];
+    const binary = atob(base64);
+    const binaryLength = binary.length;
+    const arrayBuffer = new Uint8Array(binaryLength);
+    for (let i = 0; i < binaryLength; i++) {
+      arrayBuffer[i] = binary.charCodeAt(i);
+    }
 
-  const binary = atob(base64);
-  const binaryLength = binary.length;
-  const arrayBuffer = new Uint8Array(binaryLength);
-  for (let i = 0; i < binaryLength; i++) {
-    arrayBuffer[i] = binary.charCodeAt(i);
+    const blob = new Blob([arrayBuffer], { type: mimeType });
+    const file = new File([blob], filename, { type: mimeType });
+    return file;
+  } catch (error) {
+    console.error("Error converting base64 to image:", error);
+    throw error;
   }
-
-  const blob = new Blob([arrayBuffer], { type: "image/png" });
-  const file = new File([blob], filename, { type: "image/png" });
-  return file;
 }
 
 export const openAiTools: AssistantTool[] = [