When processing the response, handle it differently if it's just a text response

Author: sabaimran

src/flint/routers/api.py

@@ -234,25 +234,26 @@ async def response_to_user_whatsapp(message: str, from_number: str, body, intro_
     chat_response = send_message_to_khoj_chat(user_message, from_number)
 
     if chat_response.get("response"):
-        chat_response_text = chat_response["response"]
-        if chat_response_text.get("image"):
-            media_url = chat_response_text["image"]
-            if media_url:
-                # Write the file to a tmp directory
-                filepath = f"/tmp/{int(time.time() * 1000)}.png"
-                response = requests.get(media_url)
-                response.raise_for_status()
-
-                # The incoming image is a link to a webp image. We need to convert it to a png image.
-                image = Image.open(BytesIO(response.content))
-                image.save(filepath, "PNG")
-
-                media_id = upload_media_to_whatsapp(filepath, "image/png", phone_number_id)
-                data = make_whatsapp_image_payload(media_id, from_number)
-                response = whatsapp_cloud_api_session.post(url, json=data)
-                response.raise_for_status()
-                os.remove(filepath)
-        else:
+        chat_response_text = chat_response.get("response")
+        try:
+            if chat_response_text.get("image"):
+                media_url = chat_response_text["image"]
+                if media_url:
+                    # Write the file to a tmp directory
+                    filepath = f"/tmp/{int(time.time() * 1000)}.png"
+                    response = requests.get(media_url)
+                    response.raise_for_status()
+
+                    # The incoming image is a link to a webp image. We need to convert it to a png image.
+                    image = Image.open(BytesIO(response.content))
+                    image.save(filepath, "PNG")
+
+                    media_id = upload_media_to_whatsapp(filepath, "image/png", phone_number_id)
+                    data = make_whatsapp_image_payload(media_id, from_number)
+                    response = whatsapp_cloud_api_session.post(url, json=data)
+                    response.raise_for_status()
+                    os.remove(filepath)
+        except AttributeError:
             data = make_whatsapp_payload(chat_response_text, from_number)
             response = whatsapp_cloud_api_session.post(url, json=data)
             response.raise_for_status()

src/flint/routers/dev.py

@@ -27,15 +27,17 @@ async def chat_dev(
 
     if chat_response.get("response"):
         response = chat_response.get("response")
-        if response.get("image"):
-            encoded_img = response["image"]
-            if encoded_img:
-                # Write the file to a tmp directory
-                filepath = f"/tmp/{int(time.time() * 1000)}.png"
-                with open(filepath, "wb") as f:
-                    f.write(base64.b64decode(encoded_img))
-            chat_response_text = f"Image saved to {filepath}"
-        else:
+        # Try parsing to see if the response is a JSON object
+        try:
+            if response.get("image"):
+                encoded_img = response["image"]
+                if encoded_img:
+                    # Write the file to a tmp directory
+                    filepath = f"/tmp/{int(time.time() * 1000)}.png"
+                    with open(filepath, "wb") as f:
+                        f.write(base64.b64decode(encoded_img))
+                chat_response_text = f"Image saved to {filepath}"
+        except AttributeError:
             chat_response_text = chat_response["response"]
     elif chat_response.get("detail"):
         chat_response_text = chat_response["detail"]