Update app.py

app.py

@@ -153,29 +153,45 @@ import re
 import mimetypes # Added for MIME type detection
 import requests
 from typing import List, Dict, Union, Any
+import base64
+import re
+import mimetypes # For fallback MIME type detection
 
 def encode_image_to_base64(image_path: str) -> str:
-    """Convert an image file to base64 encoding."""
+    """Encodes a local image file to a base64 string."""
     with open(image_path, "rb") as image_file:
         return base64.b64encode(image_file.read()).decode('utf-8')
 
-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':
+def get_image_media_type(image_path: str) -> str | None:
+    """
+    Determines the media type for an image file.
+    Returns the MIME string (e.g., "image/jpg") or None if not a recognized image type.
+    """
+    ext = image_path.split('.')[-1].lower()
+    
+    if ext in ("jpg", "jpeg"):
+        return "image/jpg"  # Align with the example snippet's "image/jpg"
+    elif ext == "png":
         return "image/png"
-    elif extension == 'gif':
+    elif ext == "gif":
         return "image/gif"
-    elif extension == 'webp':
+    elif ext == "webp":
         return "image/webp"
-    else:
-        return f"image/{extension}"
+    
+    # Fallback to mimetypes for other potential image types
+    mime_type, _ = mimetypes.guess_type(image_path)
+    if mime_type and mime_type.startswith("image/"):
+        if mime_type == "image/jpeg": # If mimetypes returns image/jpeg, use image/jpg
+            return "image/jpg"
+        return mime_type
+        
+    return None # Not a recognized/supported image type
 
-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"])
+def process_interleaved_images(message: dict) -> list:
+    """Processes messages with <image> tags interleaved with text."""
+    user_text = message.get("text", "")
+    files = message.get("files", [])
+    parts = re.split(r"(<image>)", user_text)
     
     final_content = []
     current_text = ""
@@ -183,96 +199,128 @@ def process_interleaved_images(message: Dict[str, Any]) -> List[Dict[str, Any]]:
     
     for part in parts:
         if part == "<image>":
-            # Add accumulated text if any
             if current_text.strip():
                 final_content.append({"type": "text", "text": current_text.strip()})
                 current_text = ""
             
-            # 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:{mime_type};base64,{base64_data}"}
-                })
+            if image_index < len(files):
+                image_path = files[image_index]
+                media_type = get_image_media_type(image_path)
+                if media_type: # Proceed only if it's a recognized image type
+                    try:
+                        base64_image = encode_image_to_base64(image_path)
+                        final_content.append({
+                            "type": "image_url",
+                            "image_url": {"url": f"data:{media_type};base64,{base64_image}"}
+                        })
+                    except FileNotFoundError:
+                        # Optionally log this error or add a placeholder for the missing image
+                        print(f"Warning: Image file not found: {image_path}")
+                    except Exception as e:
+                        print(f"Warning: Could not process image {image_path}: {e}")
+                else:
+                    # File is not a recognized image type, or get_image_media_type returned None
+                    print(f"Warning: File {image_path} is not a recognized image type or <image> tag mismatch.")
                 image_index += 1
+            else:
+                # More <image> tags than files provided
+                print("Warning: <image> tag found but no corresponding file path in 'files' list.")
         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_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) -> list:
+    """Processes a new user message, handling text, images, and potentially video."""
+    user_text = message.get("text", "")
+    files = message.get("files", [])
 
-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"]}]
-    
-    # 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])
+    if not files:
+        return [{"role": "user", "content": user_text}]
+
+    if files and files[0].endswith(".mp4"):
+        text_message = {"role": "user", "content": user_text}
+        video_messages = process_video(files[0]) # process_video needs to be defined
         return [text_message] + video_messages
-    
-    # Handle interleaved images
-    if "<image>" in message["text"]:
-        content = process_interleaved_images(message)
+
+    if "<image>" in user_text:
+        content = process_interleaved_images(message) # Pass the whole message dictionary
         return [{"role": "user", "content": content}]
     
-    # Handle text with images appended
-    content = [{"type": "text", "text": message["text"]}]
-    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:{mime_type};base64,{base64_data}"}
-        })
+    # For text with images appended (if no <image> tags or if files exist beyond those for tags)
+    content = []
+    if user_text.strip(): # Add text part only if there's text
+        content.append({"type": "text", "text": user_text})
+
+    for path in files:
+        # This simplistic check assumes non-mp4 files could be images.
+        # If interleaved images already consumed some files, this might re-process or process remaining.
+        # A more sophisticated approach might be needed if mixing interleaved and appended from the same 'files' list.
+        if not path.endswith(".mp4"): 
+            media_type = get_image_media_type(path)
+            if media_type: # Proceed only if it's a recognized image type
+                try:
+                    base64_image = encode_image_to_base64(path)
+                    content.append({
+                        "type": "image_url",
+                        "image_url": {"url": f"data:{media_type};base64,{base64_image}"}
+                    })
+                except FileNotFoundError:
+                    print(f"Warning: Image file not found during append: {path}")
+                except Exception as e:
+                    print(f"Warning: Could not process image {path} during append: {e}")
+            # else: file is not a recognized image, skip.
     
+    if not content: # If after processing, content is empty (e.g. only non-image files provided with no text)
+        return [{"role": "user", "content": ""}] # Send empty content rather than nothing
+
     return [{"role": "user", "content": content}]
 
-def process_history(history: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-    """Process conversation history, encoding images as base64."""
+def process_history(history: list[dict]) -> list[dict]:
+    """Processes chat history, converting file:// image URLs to base64 data URLs."""
     messages = []
     
     for item in history:
         if item["role"] == "assistant":
             messages.append({"role": "assistant", "content": item["content"]})
         else:  # user messages
-            content = item.get("content")
-            if isinstance(content, str):
-                messages.append({"role": "user", "content": content})
-            elif isinstance(content, list) and len(content) > 0 and isinstance(content[0], str):
-                # Image file path
-                image_path = content[0]
-                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:{mime_type};base64,{base64_data}"}
-                        }
-                    ]
-                })
+            current_content = item.get("content")
+            if isinstance(current_content, str):
+                messages.append({"role": "user", "content": current_content})
+            elif isinstance(current_content, list): # 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://"
+                        media_type = get_image_media_type(image_path)
+                        if media_type: # Proceed only if it's a recognized image type
+                            try:
+                                base64_image = encode_image_to_base64(image_path)
+                                processed_content_parts.append({
+                                    "type": "image_url",
+                                    "image_url": {"url": f"data:{media_type};base64,{base64_image}"}
+                                })
+                            except FileNotFoundError:
+                                print(f"Warning: History image file not found: {image_path}")
+                                processed_content_parts.append(part) # Keep original part if file missing
+                            except Exception as e:
+                                print(f"Warning: Could not process history image {image_path}: {e}")
+                                processed_content_parts.append(part) # Keep original part on other errors
+                        else:
+                            # Was a file:// URL but not a recognized image or path issue
+                            print(f"Warning: History file {image_path} is not a recognized image type.")
+                            processed_content_parts.append(part) # Keep original part
+                    else:
+                        processed_content_parts.append(part)
+                messages.append({"role": "user", "content": processed_content_parts})
             else:
-                # Already properly formatted content
-                messages.append({"role": "user", "content": content})
-    
+                # Content is not a string or list, pass as is or log warning
+                messages.append({"role": "user", "content": current_content if current_content is not None else ""})
+                
     return messages
 
 def run(message: dict, history: list[dict]) -> Iterator[str]: