Update app.py

app.py

@@ -147,14 +147,32 @@ def process_video(video_path: str) -> list[dict]:
             })
     
     return image_messages
+    
+import base64
+import re
+import mimetypes # Added for MIME type detection
 
-
-def encode_image_to_base64(image_path):
-    import mimetypes
-    mime_type, _ = mimetypes.guess_type(image_path)
+def encode_image_to_base64(image_path: str) -> str:
     with open(image_path, "rb") as image_file:
-        encoded = base64.b64encode(image_file.read()).decode("utf-8")
-    return f"data:{mime_type};base64,{encoded}"
+        return base64.b64encode(image_file.read()).decode('utf-8')
+
+def get_image_media_type(image_path: str) -> str:
+    ext = image_path.split('.')[-1].lower()
+    if ext in ("jpg", "jpeg"):
+        return "image/jpeg"
+    elif ext == "png":
+        return "image/png"
+    elif ext == "gif":
+        return "image/gif"
+    elif ext == "webp":
+        return "image/webp"
+    else:
+        # Fallback to mimetypes detection
+        mime_type, _ = mimetypes.guess_type(image_path)
+        if mime_type and mime_type.startswith("image/"):
+            return mime_type
+        # Default fallback if type is unknown or not a recognized image type
+        return "application/octet-stream"
 
 
 def process_interleaved_images(message: dict) -> list:
@@ -166,72 +184,89 @@ def process_interleaved_images(message: dict) -> list:
     
     for part in parts:
         if part == "<image>":
+            # If we have accumulated text, add it first
             if current_text.strip():
                 final_content.append({"type": "text", "text": current_text.strip()})
                 current_text = ""
-            encoded_image = encode_image_to_base64(message["files"][image_index])
-            final_content.append({
-                "type": "image_url",
-                "image_url": {"url": encoded_image}
-            })
-            image_index += 1
+            
+            # Add the image as base64 data URL
+            if image_index < len(message['files']):
+                image_path = message['files'][image_index]
+                base64_image = encode_image_to_base64(image_path)
+                media_type = get_image_media_type(image_path)
+                final_content.append({
+                    "type": "image_url",
+                    "image_url": {"url": f"data:{media_type};base64,{base64_image}"}
+                })
+                image_index += 1
         else:
             current_text += part
-
+    
+    # Add any remaining text
     if current_text.strip():
         final_content.append({"type": "text", "text": current_text.strip()})
-
+    
     return final_content
 
 
-def process_new_user_message(message: dict):
-    if not message["files"]:
+def process_new_user_message(message: dict) -> list:
+    if not message.get("files"): # Check if "files" key exists and is not empty
         return [{"role": "user", "content": message["text"]}]
 
     if message["files"][0].endswith(".mp4"):
         text_message = {"role": "user", "content": message["text"]}
-        video_messages = process_video(message["files"][0])
+        video_messages = process_video(message["files"][0]) # process_video needs to be defined
         return [text_message] + video_messages
 
     if "<image>" in message["text"]:
+        # For interleaved text and images
         content = process_interleaved_images(message)
         return [{"role": "user", "content": content}]
     
     # For text with images appended
     content = [{"type": "text", "text": message["text"]}]
     for path in message["files"]:
-        encoded_image = encode_image_to_base64(path)
-        content.append({
-            "type": "image_url",
-            "image_url": {"url": encoded_image}
-        })
+        if not path.endswith(".mp4"): # Simple check to avoid processing videos as images
+            base64_image = encode_image_to_base64(path)
+            media_type = get_image_media_type(path)
+            content.append({
+                "type": "image_url",
+                "image_url": {"url": f"data:{media_type};base64,{base64_image}"}
+            })
     
     return [{"role": "user", "content": content}]
 
 
 def process_history(history: list[dict]) -> list[dict]:
     messages = []
-
+    
     for item in history:
         if item["role"] == "assistant":
             messages.append({"role": "assistant", "content": item["content"]})
-        else:
-            content = item["content"]
-            if isinstance(content, str):
-                messages.append({"role": "user", "content": content})
-            else:
-                # Assume content[0] is a file path
-                encoded_image = encode_image_to_base64(content[0])
-                messages.append({
-                    "role": "user", 
-                    "content": [
-                        {
-                            "type": "image_url",
-                            "image_url": {"url": encoded_image}
-                        }
-                    ]
-                })
-
+        else:  # user messages
+            current_content = item["content"]
+            if isinstance(current_content, str):
+                messages.append({"role": "user", "content": current_content})
+            elif isinstance(current_content, list): # Handles multimodal content (list of dicts)
+                processed_content_parts = []
+                for part in current_content:
+                    if part.get("type") == "image_url" and \
+                       part.get("image_url", {}).get("url", "").startswith("file://"):
+                        image_path = part["image_url"]["url"][7:] # Remove "file://"
+                        try:
+                            base64_image = encode_image_to_base64(image_path)
+                            media_type = get_image_media_type(image_path)
+                            processed_content_parts.append({
+                                "type": "image_url",
+                                "image_url": {"url": f"data:{media_type};base64,{base64_image}"}
+                            })
+                        except FileNotFoundError:
+                            # Handle missing file if necessary, e.g., skip or keep original
+                            processed_content_parts.append(part) # Keep original if file not found
+                    else:
+                        processed_content_parts.append(part)
+                messages.append({"role": "user", "content": processed_content_parts})
+    
     return messages