Update app.py

app.py

@@ -152,30 +152,13 @@ import base64
 import re
 import mimetypes # Added for MIME type detection
 
-def encode_image_to_base64(image_path: str) -> str:
+def encode_image_to_base64(image_path):
+    import base64
     with open(image_path, "rb") as image_file:
         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:
+    import re
     parts = re.split(r"(<image>)", message["text"])
     
     final_content = []
@@ -189,16 +172,20 @@ def process_interleaved_images(message: dict) -> list:
                 final_content.append({"type": "text", "text": current_text.strip()})
                 current_text = ""
             
-            # 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
+            # Add the image as base64
+            image_path = message['files'][image_index]
+            image_base64 = encode_image_to_base64(image_path)
+            # Determine media type based on file extension
+            image_extension = image_path.split('.')[-1].lower()
+            media_type = f"image/{image_extension}"
+            if image_extension == 'jpg':
+                media_type = "image/jpeg"
+                
+            final_content.append({
+                "type": "image_url",
+                "image_url": {"url": f"data:{media_type};base64,{image_base64}"}
+            })
+            image_index += 1
         else:
             current_text += part
     
@@ -208,16 +195,14 @@ def process_interleaved_images(message: dict) -> list:
     
     return final_content
 
-
-def process_new_user_message(message: dict) -> list:
-    if not message.get("files"): # Check if "files" key exists and is not empty
+def process_new_user_message(message: dict):
+    if not message["files"]:
         return [{"role": "user", "content": message["text"]}]
-
     if message["files"][0].endswith(".mp4"):
+        # For video, return text message followed by frame messages
         text_message = {"role": "user", "content": message["text"]}
-        video_messages = process_video(message["files"][0]) # process_video needs to be defined
+        video_messages = process_video(message["files"][0])
         return [text_message] + video_messages
-
     if "<image>" in message["text"]:
         # For interleaved text and images
         content = process_interleaved_images(message)
@@ -226,17 +211,21 @@ def process_new_user_message(message: dict) -> list:
     # For text with images appended
     content = [{"type": "text", "text": message["text"]}]
     for path in message["files"]:
-        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}"}
-            })
+        # Convert image to base64
+        image_base64 = encode_image_to_base64(path)
+        # Determine media type based on file extension
+        image_extension = path.split('.')[-1].lower()
+        media_type = f"image/{image_extension}"
+        if image_extension == 'jpg':
+            media_type = "image/jpeg"
+            
+        content.append({
+            "type": "image_url",
+            "image_url": {"url": f"data:{media_type};base64,{image_base64}"}
+        })
     
     return [{"role": "user", "content": content}]
 
-
 def process_history(history: list[dict]) -> list[dict]:
     messages = []
     
@@ -244,32 +233,31 @@ def process_history(history: list[dict]) -> list[dict]:
         if item["role"] == "assistant":
             messages.append({"role": "assistant", "content": item["content"]})
         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"]
-                        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})
-    print(messages)
+            content = item["content"]
+            if isinstance(content, str):
+                messages.append({"role": "user", "content": content})
+            else:  # image content
+                # Convert image to base64
+                image_path = content[0]
+                image_base64 = encode_image_to_base64(image_path)
+                # Determine media type based on file extension
+                image_extension = image_path.split('.')[-1].lower()
+                media_type = f"image/{image_extension}"
+                if image_extension == 'jpg':
+                    media_type = "image/jpeg"
+                    
+                messages.append({
+                    "role": "user", 
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": f"data:{media_type};base64,{image_base64}"}
+                        }
+                    ]
+                })
+    
     return messages
 
-
 def run(message: dict, history: list[dict]) -> Iterator[str]:
     if not validate_media_constraints(message, history):
         yield ""