Update app.py

app.py

@@ -151,14 +151,33 @@ def process_video(video_path: str) -> list[dict]:
 import base64
 import re
 import mimetypes # Added for MIME type detection
+import base64
+import re
+import os
+import requests
+from typing import List, Dict, Union, Any
 
-def encode_image_to_base64(image_path):
-    import base64
+def encode_image_to_base64(image_path: str) -> str:
+    """Convert an image file to base64 encoding."""
     with open(image_path, "rb") as image_file:
         return base64.b64encode(image_file.read()).decode('utf-8')
 
-def process_interleaved_images(message: dict) -> list:
-    import re
+def get_mime_type(file_path: str) -> str:
+    """Determine MIME type based on file extension."""
+    extension = file_path.split('.')[-1].lower()
+    if extension == 'jpg' or extension == 'jpeg':
+        return "image/jpeg"
+    elif extension == 'png':
+        return "image/png"
+    elif extension == 'gif':
+        return "image/gif"
+    elif extension == 'webp':
+        return "image/webp"
+    else:
+        return f"image/{extension}"
+
+def process_interleaved_images(message: Dict[str, Any]) -> List[Dict[str, Any]]:
+    """Process a message with interleaved text and image tags."""
     parts = re.split(r"(<image>)", message["text"])
     
     final_content = []
@@ -167,25 +186,22 @@ def process_interleaved_images(message: dict) -> list:
     
     for part in parts:
         if part == "<image>":
-            # If we have accumulated text, add it first
+            # Add accumulated text if any
             if current_text.strip():
                 final_content.append({"type": "text", "text": current_text.strip()})
                 current_text = ""
             
-            # 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"
+            # Process and add the image as base64
+            if image_index < len(message.get('files', [])):
+                image_path = message['files'][image_index]
+                base64_data = encode_image_to_base64(image_path)
+                mime_type = get_mime_type(image_path)
                 
-            final_content.append({
-                "type": "image_url",
-                "image_url": {"url": f"data:{media_type};base64,{image_base64}"}
-            })
-            image_index += 1
+                final_content.append({
+                    "type": "image_url",
+                    "image_url": {"url": f"data:{mime_type};base64,{base64_data}"}
+                })
+                image_index += 1
         else:
             current_text += part
     
@@ -195,66 +211,70 @@ def process_interleaved_images(message: dict) -> list:
     
     return final_content
 
-def process_new_user_message(message: dict):
-    if not message["files"]:
+def process_video(video_path: str) -> List[Dict[str, Any]]:
+    """Placeholder for video processing function."""
+    # Implementation depends on requirements for video processing
+    pass
+
+def process_new_user_message(message: Dict[str, Any]) -> List[Dict[str, Any]]:
+    """Process user message with potential images or videos."""
+    # If no files, just return the text
+    if not message.get("files", []):
         return [{"role": "user", "content": message["text"]}]
-    if message["files"][0].endswith(".mp4"):
-        # For video, return text message followed by frame messages
+    
+    # Handle video files
+    if message["files"] and message["files"][0].endswith(".mp4"):
         text_message = {"role": "user", "content": message["text"]}
         video_messages = process_video(message["files"][0])
         return [text_message] + video_messages
+    
+    # Handle interleaved images
     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
+    # Handle text with images appended
     content = [{"type": "text", "text": message["text"]}]
-    for path in message["files"]:
-        # 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"
-            
+    for file_path in message.get("files", []):
+        base64_data = encode_image_to_base64(file_path)
+        mime_type = get_mime_type(file_path)
+        
         content.append({
             "type": "image_url",
-            "image_url": {"url": f"data:{media_type};base64,{image_base64}"}
+            "image_url": {"url": f"data:{mime_type};base64,{base64_data}"}
         })
     
     return [{"role": "user", "content": content}]
 
-def process_history(history: list[dict]) -> list[dict]:
+def process_history(history: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Process conversation history, encoding images as base64."""
     messages = []
     
     for item in history:
         if item["role"] == "assistant":
             messages.append({"role": "assistant", "content": item["content"]})
         else:  # user messages
-            content = item["content"]
+            content = item.get("content")
             if isinstance(content, str):
                 messages.append({"role": "user", "content": content})
-            else:  # image content
-                # Convert image to base64
+            elif isinstance(content, list) and len(content) > 0 and isinstance(content[0], str):
+                # Image file path
                 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"
-                    
+                base64_data = encode_image_to_base64(image_path)
+                mime_type = get_mime_type(image_path)
+                
                 messages.append({
                     "role": "user", 
                     "content": [
                         {
                             "type": "image_url",
-                            "image_url": {"url": f"data:{media_type};base64,{image_base64}"}
+                            "image_url": {"url": f"data:{mime_type};base64,{base64_data}"}
                         }
                     ]
                 })
+            else:
+                # Already properly formatted content
+                messages.append({"role": "user", "content": content})
     
     return messages