From c69edeb4b6c954d12b25facb9c2b61e6c0f52027 Mon Sep 17 00:00:00 2001
From: Abi Raja <abimanyuraja@gmail.com>
Date: Tue, 5 Mar 2024 14:02:48 -0500
Subject: [PATCH] send correct media type to Claude

---
 backend/llm.py | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/backend/llm.py b/backend/llm.py
index c6d629e..f042223 100644
--- a/backend/llm.py
+++ b/backend/llm.py
@@ -69,13 +69,23 @@ async def stream_claude_response(
         for content in message["content"]:  # type: ignore
             if content["type"] == "image_url":
                 content["type"] = "image"
-                content["source"] = {
-                    "type": "base64",
-                    "media_type": "image/png",  # TODO: Automatically detect media type
-                    "data": content["image_url"]["url"].split(",")[1],  # type: ignore
-                }
+
+                # Extract base64 data and media type from data URL
+                # Example base64 data URL: data:image/png;base64,iVBOR...
+                image_data_url = cast(str, content["image_url"]["url"])
+                media_type = image_data_url.split(";")[0].split(":")[1]
+                base64_data = image_data_url.split(",")[1]
+
+                # Remove OpenAI parameter
                 del content["image_url"]
 
+                content["source"] = {
+                    "type": "base64",
+                    "media_type": media_type,
+                    "data": base64_data,
+                }
+
+    # Stream Claude response
     async with client.messages.stream(
         model=model,
         max_tokens=max_tokens,
@@ -86,5 +96,6 @@ async def stream_claude_response(
         async for text in stream.text_stream:
             await callback(text)
 
+    # Return final message
     response = await stream.get_final_message()
     return response.content[0].text