Update app.py

app.py

@@ -219,62 +219,87 @@ def upload_file_to_mistral(client, path, filename=None, purpose="batch"):
 
 
 def build_messages_for_image(prompt: str, b64_jpg: str = None, image_url: str = None):
-    user_content = [{"type": "text", "text": prompt}]
+    # Keep user-visible prompt and include a single string content containing a data URL or image URL.
     if image_url:
-        user_content.append({"type": "image_url", "image_url": image_url})
+        content = f"{prompt}\n\nImage: {image_url}"
     elif b64_jpg:
-        user_content.append({"type": "image_base64", "image_base64": b64_jpg})
+        # Use explicit data URL so Mistral chat content is a single string (fixes Pydantic validation)
+        content = f"{prompt}\n\nImage (base64): data:image/jpeg;base64,{b64_jpg}"
     else:
         raise ValueError("Either image_url or b64_jpg required")
-    return [{"role": "system", "content": SYSTEM_INSTRUCTION}, {"role": "user", "content": user_content}]
+    return [{"role": "system", "content": SYSTEM_INSTRUCTION}, {"role": "user", "content": content}]
 
 
 def build_messages_for_text(prompt: str, extra_text: str):
     return [{"role": "system", "content": SYSTEM_INSTRUCTION}, {"role": "user", "content": f"{prompt}\n\n{extra_text}"}]
 
 
-# New helper: normalize messages so content is always a plain string
-def normalize_messages(messages):
-    out = []
-    for m in messages:
-        if not isinstance(m, dict):
-            out.append(m)
-            continue
-        c = m.get("content")
-        if isinstance(c, list):
-            parts = []
-            for item in c:
-                if isinstance(item, str):
-                    parts.append(item)
-                elif isinstance(item, dict):
-                    typ = item.get("type")
-                    if typ == "text" and item.get("text"):
-                        parts.append(item["text"])
-                    elif typ == "image_url" and item.get("image_url"):
-                        parts.append(item["image_url"])
-                    elif typ == "image_base64" and item.get("image_base64"):
-                        # convert to data URL to satisfy string requirement
-                        parts.append("data:image/jpeg;base64," + item["image_base64"])
-                    else:
-                        parts.append(item.get("text") or item.get("image_url") or item.get("image_base64") or "")
+def extract_text_from_response(res, parts: list):
+    try:
+        choices = getattr(res, "choices", None) or res.get("choices", [])
+    except Exception:
+        choices = []
+    if choices:
+        try:
+            msg = choices[0].message
+            if isinstance(msg, dict):
+                content = msg.get("content")
+            else:
+                content = getattr(msg, "content", None)
+            if content:
+                if isinstance(content, str):
+                    parts.append(content)
                 else:
-                    parts.append(str(item))
-            newc = "\n\n".join(p for p in parts if p).strip()
-            nm = m.copy()
-            nm["content"] = newc
-            out.append(nm)
-        elif not isinstance(c, str):
-            nm = m.copy()
-            nm["content"] = str(c or "")
-            out.append(nm)
-        else:
-            out.append(m)
-    return out
+                    if isinstance(content, list):
+                        for c in content:
+                            if isinstance(c, dict) and c.get("type") == "text":
+                                parts.append(c.get("text", ""))
+                    elif isinstance(content, dict):
+                        text = content.get("text") or content.get("content")
+                        if text:
+                            parts.append(text)
+        except Exception:
+            parts.append(str(res))
+    else:
+        parts.append(str(res))
 
 
 def stream_and_collect(client, model, messages, parts: list):
     try:
-        norm_msgs = normalize_messages(messages)
+        # We assume messages are plain strings already for image/text functions above.
+        # For safety, if content is a list/dict (old codepath), coerce minimally here:
+        norm_msgs = []
+        for m in messages:
+            if not isinstance(m, dict):
+                norm_msgs.append(m)
+                continue
+            c = m.get("content")
+            if isinstance(c, list):
+                # prefer image_url or image_base64 if present, else join text chunks
+                picked = []
+                for item in c:
+                    if isinstance(item, dict):
+                        if item.get("type") == "image_url" and item.get("image_url"):
+                            picked.append(item["image_url"])
+                        elif item.get("type") == "image_base64" and item.get("image_base64"):
+                            picked.append("data:image/jpeg;base64," + item["image_base64"])
+                        elif item.get("type") == "text" and item.get("text"):
+                            picked.append(item["text"])
+                    elif isinstance(item, str):
+                        picked.append(item)
+                newc = "\n\n".join(p for p in picked if p).strip()
+                nm = m.copy()
+                nm["content"] = newc
+                norm_msgs.append(nm)
+            else:
+                if not isinstance(c, str):
+                    nm = m.copy()
+                    nm["content"] = str(c or "")
+                    norm_msgs.append(nm)
+                else:
+                    norm_msgs.append(m)
+
+        # Try streaming first (client may expose .chat.stream); fall back to non-streaming complete
         stream_gen = None
         try:
             stream_gen = client.chat.stream(model=model, messages=norm_msgs)
@@ -289,34 +314,9 @@ def stream_and_collect(client, model, messages, parts: list):
                     continue
                 parts.append(d)
             return
+
         res = client.chat.complete(model=model, messages=norm_msgs, stream=False)
-        try:
-            choices = getattr(res, "choices", None) or res.get("choices", [])
-        except Exception:
-            choices = []
-        if choices:
-            try:
-                msg = choices[0].message
-                if isinstance(msg, dict):
-                    content = msg.get("content")
-                else:
-                    content = getattr(msg, "content", None)
-                if content:
-                    if isinstance(content, str):
-                        parts.append(content)
-                    else:
-                        if isinstance(content, list):
-                            for c in content:
-                                if isinstance(c, dict) and c.get("type") == "text":
-                                    parts.append(c.get("text", ""))
-                        elif isinstance(content, dict):
-                            text = content.get("text") or content.get("content")
-                            if text:
-                                parts.append(text)
-            except Exception:
-                parts.append(str(res))
-        else:
-            parts.append(str(res))
+        extract_text_from_response(res, parts)
     except Exception as e:
         parts.append(f"[Model error: {e}]")
 
@@ -328,6 +328,7 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
     is_image = ext in IMAGE_EXTS or (not is_remote(src) and os.path.isfile(src) and ext in IMAGE_EXTS)
     parts = []
 
+    # IMAGE path: keep exactly the previous behavior but send content as a single string data URL
     if is_image:
         try:
             if is_remote(src):
@@ -342,6 +343,7 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
         stream_and_collect(client, DEFAULT_IMAGE_MODEL, msgs, parts)
         return "".join(parts).strip()
 
+    # REMOTE VIDEO path
     if is_remote(src):
         try:
             media_bytes = fetch_bytes(src, timeout=120)
@@ -350,12 +352,21 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
         ext = ext_from_src(src) or ".mp4"
         tmp_media = save_bytes_to_temp(media_bytes, suffix=ext)
         try:
+            # Try uploading full video to Mistral files first
             try:
                 file_id = upload_file_to_mistral(client, tmp_media, filename=os.path.basename(src.split("?")[0]))
-            except Exception as e:
+                extra = (
+                    f"Remote video uploaded to Mistral Files with id: {file_id}\n\n"
+                    "Instruction: Analyze the video contents using the uploaded file id. Do not invent frames not present."
+                )
+                msgs = build_messages_for_text(prompt, extra)
+                stream_and_collect(client, DEFAULT_VIDEO_MODEL, msgs, parts)
+                return "".join(parts).strip()
+            except Exception as e_upload:
+                # Fallback: extract a representative frame and analyze as image
                 frame_bytes = extract_best_frame_bytes(tmp_media)
                 if not frame_bytes:
-                    return f"Error uploading to Mistral and no frame fallback available: {e}"
+                    return f"Error uploading to Mistral and no frame fallback available: {e_upload}"
                 try:
                     jpg = convert_to_jpeg_bytes(frame_bytes, base_h=480)
                 except UnidentifiedImageError:
@@ -364,11 +375,6 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
                 msgs = build_messages_for_image(prompt, b64_jpg=b64)
                 stream_and_collect(client, DEFAULT_VIDEO_MODEL, msgs, parts)
                 return "".join(parts).strip()
-
-            extra = f"Remote video uploaded to Mistral Files with id: {file_id}\n\nInstruction: Analyze the video contents using the uploaded file id. Do not invent frames not present."
-            msgs = build_messages_for_text(prompt, extra)
-            stream_and_collect(client, DEFAULT_VIDEO_MODEL, msgs, parts)
-            return "".join(parts).strip()
         finally:
             try:
                 if tmp_media and os.path.exists(tmp_media):
@@ -376,6 +382,7 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
             except Exception:
                 pass
 
+    # LOCAL VIDEO path
     tmp_media = None
     try:
         media_bytes = fetch_bytes(src)
@@ -383,12 +390,17 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
         ext = ext or ".mp4"
         tmp_media = save_bytes_to_temp(media_bytes, suffix=ext)
         try:
+            # Upload local video file to Mistral files and analyze by file id
             file_id = upload_file_to_mistral(client, tmp_media, filename=os.path.basename(src))
-            extra = f"Local video uploaded to Mistral Files with id: {file_id}\n\nInstruction: Analyze the video contents using the uploaded file id. Do not invent frames not present."
+            extra = (
+                f"Local video uploaded to Mistral Files with id: {file_id}\n\n"
+                "Instruction: Analyze the video contents using the uploaded file id. Do not invent frames not present."
+            )
             msgs = build_messages_for_text(prompt, extra)
             stream_and_collect(client, DEFAULT_VIDEO_MODEL, msgs, parts)
             return "".join(parts).strip()
         except Exception:
+            # fallback to frame extraction + image analysis
             frame_bytes = extract_best_frame_bytes(tmp_media)
             if not frame_bytes:
                 return "Unable to process the provided file. Provide a direct image/frame URL or a remote video URL."
@@ -409,12 +421,14 @@ css = ".preview_media img, .preview_media video { max-width: 100%; height: auto;
 
 
 def load_preview(url: str):
+    # Return (image_or_None, video_or_None, status_text)
     if not url:
         return None, None, ""
     try:
         r = requests.get(url, timeout=30, stream=True)
         r.raise_for_status()
         ctype = (r.headers.get("content-type") or "").lower()
+        # If content-type indicates video or URL ends with a video ext, return video preview
         if (ctype and ctype.startswith("video/")) or any(url.lower().split("?")[0].endswith(ext) for ext in VIDEO_EXTS):
             return None, url, "Video"
         data = r.content
@@ -438,16 +452,4 @@ with gr.Blocks(title="Flux", css=css) as demo:
             with gr.Accordion("Mistral API Key (optional)", open=False):
                 api_key = gr.Textbox(label="API Key", type="password", max_lines=1)
             submit = gr.Button("Submit")
-            preview_image = gr.Image(label="Preview", type="pil", elem_classes="preview_media")
-            preview_video = gr.Video(label="Preview", elem_classes="preview_media")
-
-        with gr.Column(scale=2):
-            final_text = gr.Markdown(value="")
-
-    # Ensure preview outputs get None when not applicable so Gradio hides them
-    url_input.change(fn=load_preview, inputs=[url_input], outputs=[preview_image, preview_video, gr.Textbox(visible=False)])
-    submit.click(fn=generate_final_text, inputs=[url_input, custom_prompt, api_key], outputs=[final_text])
-    demo.queue()
-
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))
+            preview_image = gr.Image(label="Preview image", type="pil", visible