Update app.py

app.py

@@ -9,8 +9,6 @@ from typing import List, Tuple, Optional
 import requests
 from PIL import Image, ImageFile, UnidentifiedImageError
 import gradio as gr
-# import asyncio # Not used
-# import threading # Not directly used for core logic, implicit in Gradio's max_threads
 import time
 import atexit
 from requests.exceptions import RequestException, HTTPError # Import for rate limiting
@@ -19,7 +17,7 @@ from requests.exceptions import RequestException, HTTPError # Import for rate li
 DEFAULT_KEY = os.getenv("MISTRAL_API_KEY", "")
 PIXTRAL_MODEL = "pixtral-12b-2409"
 VIDEO_MODEL = "voxtral-mini-latest"
-STREAM_THRESHOLD = 20 * 1024 * 1024
+STREAM_THRESHOLD = 20 * 1024 * 1024 # 20MB
 FFMPEG_BIN = shutil.which("ffmpeg")
 IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp", ".gif")
 VIDEO_EXTS = (".mp4", ".mov", ".webm", ".mkv", ".avi", ".flv")
@@ -49,10 +47,11 @@ except ImportError:
         def __init__(self, message, status_code=None, *args, **kwargs):
             super().__init__(message, *args)
             self.status_code = status_code
+            self._message = message # Store message explicitly for property
+
         @property
         def message(self):
-            return self.args[0] if self.args else "An unknown Mistral API error occurred (mistralai library not installed)."
-
+            return self._message # Return the stored message
 
 # --- Temporary File Cleanup ---
 _temp_preview_files_to_delete = []
@@ -75,48 +74,50 @@ def get_client(key: Optional[str] = None):
     Returns a Mistral client instance. If the mistralai library is not installed
     or the API key is missing, a mock client is returned which raises
     MistralAPIException with specific messages.
+    A `is_real_client` flag is added for `chat_complete`/`upload_file_to_mistral`
+    to distinguish between a real client and a dummy.
     """
     api_key = (key or "").strip() or DEFAULT_KEY
 
+    # Helper class for dummy chat/files service
+    class DummyMistralService:
+        def __init__(self, error_message: str, status_code: Optional[int] = None):
+            self._error_message = error_message
+            self._status_code = status_code
+
+        def complete(self, **kwargs):
+            raise MistralAPIException(self._error_message, status_code=self._status_code)
+
+        def upload(self, **kwargs):
+            raise MistralAPIException(self._error_message, status_code=self._status_code)
+
+    # Dummy client wrapper
+    class DummyClient:
+        def __init__(self, k, library_missing_msg: Optional[str] = None, api_key_missing_msg: Optional[str] = None):
+            self.api_key = k
+            self.is_real_client = False
+
+            # Determine which dummy service to provide
+            if not k: # API key missing
+                error_msg = api_key_missing_msg or "Mistral API key is not set."
+                status_code = 401
+            else: # Library missing (and key is present)
+                error_msg = library_missing_msg or "Mistral client library is not installed."
+                status_code = 500
+
+            self.chat = DummyMistralService(error_msg, status_code)
+            self.files = DummyMistralService(error_msg, status_code)
+
     if not api_key:
-        class Dummy:
-            def __init__(self, k): self.api_key = k
-            class Chat:
-                def complete(self, **kwargs):
-                    raise MistralAPIException(
-                        "Mistral API key is not set. Please provide it in the UI or as MISTRAL_API_KEY environment variable.",
-                        status_code=401
-                    )
-            chat = Chat()
-            class Files:
-                def upload(self, **kwargs):
-                    raise MistralAPIException(
-                        "Mistral API key is not set. Please provide it in the UI or as MISTRAL_API_KEY environment variable.",
-                        status_code=401
-                    )
-            files = Files()
-        return Dummy(api_key)
+        return DummyClient(api_key, api_key_missing_msg="Mistral API key is not set. Please provide it in the UI or as MISTRAL_API_KEY environment variable.")
 
     if Mistral is None:
-        class Dummy:
-            def __init__(self, k): self.api_key = k
-            class Chat:
-                def complete(self, **kwargs):
-                    raise MistralAPIException(
-                        "Mistral client library is not installed. Please install it with 'pip install mistralai'.",
-                        status_code=500
-                    )
-            chat = Chat()
-            class Files:
-                def upload(self, **kwargs):
-                    raise MistralAPIException(
-                        "Mistral client library is not installed. Please install it with 'pip install mistralai'.",
-                        status_code=500
-                    )
-            files = Files()
-        return Dummy(api_key)
+        return DummyClient(api_key, library_missing_msg="Mistral client library is not installed. Please install it with 'pip install mistralai'.")
 
-    return Mistral(api_key=api_key)
+    # If Mistral is installed and API key is present, return real client
+    client = Mistral(api_key=api_key)
+    setattr(client, 'is_real_client', True) # Mark as real client
+    return client
 
 def is_remote(src: str) -> bool:
     return bool(src) and src.startswith(("http://", "https://"))
@@ -264,18 +265,19 @@ def _get_video_info_and_timestamps(media_path: str, sample_count: int) -> Tuple[
     timestamps: List[float] = []
     if duration > 0 and sample_count > 0:
         # Ensure we don't go past the video duration, or request too many samples
-        actual_sample_count = min(sample_count, int(duration)) # Limit samples to integer seconds
+        # Limit samples to integer seconds to avoid issues with very short videos and excessive samples
+        actual_sample_count = min(sample_count, max(1, int(duration)))
         if actual_sample_count > 0:
             step = duration / (actual_sample_count + 1)
             timestamps = [step * (i + 1) for i in range(actual_sample_count)]
 
     # If no timestamps generated or duration is 0, fall back to fixed times
     if not timestamps:
-        timestamps = [0.5, 1.0, 2.0, 3.0, 4.0][:sample_count] 
+        timestamps = [0.5, 1.0, 2.0, 3.0, 4.0][:sample_count] # Use fixed initial seconds
 
     return info, timestamps
 
-def extract_frames_for_model_and_gallery(media_path: str, sample_count: int = 5, timeout_extract: int = 15, gallery_base_h: int = 720, progress=None) -> Tuple[List[bytes], List[str]]:
+def extract_frames_for_model_and_gallery(media_path: str, sample_count: int = 5, timeout_extract: int = 15, gallery_base_h: int = 1080, model_base_h: int = 1024, progress=None) -> Tuple[List[bytes], List[str]]:
     """
     Extracts frames from a video, processes them for both model input (high-res JPEG bytes)
     and gallery display (smaller JPEG temp file paths), in a single pass.
@@ -304,6 +306,7 @@ def extract_frames_for_model_and_gallery(media_path: str, sample_count: int = 5,
         fd_raw, tmp_png_path = tempfile.mkstemp(suffix=f"_frame_{i}.png")
         os.close(fd_raw)
 
+        # ffmpeg command to extract a single frame
         cmd_extract = [
             FFMPEG_BIN, "-nostdin", "-y", "-ss", str(t), "-i", media_path,
             "-frames:v", "1", "-pix_fmt", "rgb24", tmp_png_path,
@@ -317,14 +320,14 @@ def extract_frames_for_model_and_gallery(media_path: str, sample_count: int = 5,
                     raw_frame_bytes = f.read()
 
                 # For model: convert to high-res JPEG bytes (model expects this)
-                jpeg_model_bytes = convert_to_jpeg_bytes(raw_frame_bytes, base_h=720) # Keep higher res for model
+                jpeg_model_bytes = convert_to_jpeg_bytes(raw_frame_bytes, base_h=model_base_h) # Use model_base_h
                 if jpeg_model_bytes: # Only append if conversion was successful
                     frames_for_model.append(jpeg_model_bytes)
                 else:
                     print(f"Warning: Failed to convert extracted frame {i+1} to JPEG for model input.")
 
                 # For gallery: convert to smaller JPEG bytes and save as new temp file
-                jpeg_gallery_bytes = convert_to_jpeg_bytes(raw_frame_bytes, base_h=gallery_base_h) # Higher base_h
+                jpeg_gallery_bytes = convert_to_jpeg_bytes(raw_frame_bytes, base_h=gallery_base_h) # Use gallery_base_h
                 if jpeg_gallery_bytes: # Only create temp file if conversion was successful
                     temp_jpeg_path = _temp_file(jpeg_gallery_bytes, suffix=f"_gallery_{i}.jpg") # _temp_file tracks this for cleanup
                     if temp_jpeg_path: # Only add to gallery if temp file was successfully created
@@ -353,23 +356,28 @@ def chat_complete(client, model: str, messages, timeout: int = 120, progress=Non
             if progress is not None:
                 progress(0.6 + 0.01 * attempt, desc=f"Sending request to model (attempt {attempt+1}/{max_retries})...")
 
-            # Prefer using the Mistral client if available and functional
-            if hasattr(client, "chat") and hasattr(client.chat, "complete"):
+            # Check if it's the real Mistral client or a dummy/fallback to direct requests
+            if getattr(client, 'is_real_client', False):
                 res = client.chat.complete(model=model, messages=messages, stream=False, timeout_ms=timeout * 1000)
-            else: # This path should ideally not be reached if get_client returns a dummy with exceptions
-                api_key = getattr(client, "api_key", "") or DEFAULT_KEY
+                choices = getattr(res, "choices", None)
+            else: # Fallback to direct requests if Mistral library not installed or API key missing
+                api_key = getattr(client, "api_key", "") # Get API key from the client object (real or dummy)
                 if not api_key:
-                    return "Error: Mistral API key is not set. Cannot make direct requests."
+                    # This case should ideally be caught by get_client returning a DummyClient
+                    # but provides a safety net if somehow we bypass that.
+                    raise MistralAPIException(
+                        "Mistral API key is not set for direct HTTP request.", status_code=401
+                    )
                 url = "https://api.mistral.ai/v1/chat/completions"
                 headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
                 r = requests.post(url, json={"model": model, "messages": messages}, headers=headers, timeout=timeout)
                 r.raise_for_status() # Raises HTTPError for bad responses (4xx or 5xx)
                 res = r.json()
+                choices = res.get("choices") # For dictionary response from direct requests
 
             if progress is not None:
                 progress(0.8, desc="Model responded, parsing...")
 
-            choices = getattr(res, "choices", None) or (res.get("choices") if isinstance(res, dict) else [])
             if not choices:
                 return f"Empty response from model: {res}"
 
@@ -413,7 +421,8 @@ def upload_file_to_mistral(client, path: str, filename: str | None = None, purpo
             if progress is not None:
                 progress(0.5 + 0.01 * attempt, desc=f"Uploading file to model service (attempt {attempt+1}/{max_retries})...")
 
-            if hasattr(client, "files") and hasattr(client.files, "upload"):
+            # Check if it's the real Mistral client or a dummy/fallback to direct requests
+            if getattr(client, 'is_real_client', False):
                 with open(path, "rb") as fh:
                     res = client.files.upload(file={"file_name": fname, "content": fh}, purpose=purpose)
                 fid = getattr(res, "id", None) or (res.get("id") if isinstance(res, dict) else None)
@@ -424,10 +433,12 @@ def upload_file_to_mistral(client, path: str, filename: str | None = None, purpo
                 if progress is not None:
                     progress(0.6, desc="Upload complete")
                 return fid
-            else: # This path should ideally not be reached if get_client returns a dummy with exceptions
-                api_key = getattr(client, "api_key", "") or DEFAULT_KEY
+            else: # Fallback to direct requests if Mistral library not installed or API key missing
+                api_key = getattr(client, "api_key", "")
                 if not api_key:
-                    raise RuntimeError("Mistral API key is not set for file upload.")
+                    raise MistralAPIException(
+                        "Mistral API key is not set for file upload.", status_code=401
+                    )
                 url = "https://api.mistral.ai/v1/files"
                 headers = {"Authorization": f"Bearer {api_key}"}
                 with open(path, "rb") as fh:
@@ -491,7 +502,7 @@ def analyze_image_structured(client, img_bytes: bytes, prompt: str, progress=Non
     try:
         if progress is not None:
             progress(0.3, desc="Preparing image for analysis...")
-        jpeg = convert_to_jpeg_bytes(img_bytes, base_h=1024)
+        jpeg = convert_to_jpeg_bytes(img_bytes, base_h=1024) # High-res for model input
         if not jpeg: # Handle case where convert_to_jpeg_bytes failed
             return "Error: Could not convert image for analysis."
         data_url = b64_bytes(jpeg, mime="image/jpeg")
@@ -528,9 +539,9 @@ def analyze_video_cohesive(client, video_path: str, prompt: str, progress=None)
         ]
         result = chat_complete(client, VIDEO_MODEL, messages, progress=progress)
         # If successful upload, still extract frames for gallery display
-        # Use the combined function for gallery frames
+        # Use the combined function for gallery frames, now with higher resolution
         _, gallery_frame_paths = extract_frames_for_model_and_gallery(
-            video_path, sample_count=6, gallery_base_h=720, progress=progress # Increased gallery_base_h
+            video_path, sample_count=6, gallery_base_h=1080, model_base_h=1024, progress=progress
         )
         return result, gallery_frame_paths
     except Exception as e:
@@ -538,8 +549,9 @@ def analyze_video_cohesive(client, video_path: str, prompt: str, progress=None)
             progress(0.35, desc=f"Upload failed for video ({e}). Extracting frames as fallback...")
 
         # Use the combined extraction function for both model input and gallery display
+        # Higher resolution for both model frames and gallery frames
         frames_for_model_bytes, gallery_frame_paths = extract_frames_for_model_and_gallery(
-            video_path, sample_count=6, gallery_base_h=720, progress=progress # Increased gallery_base_h
+            video_path, sample_count=6, gallery_base_h=1080, model_base_h=1024, progress=progress
         )
 
         if not frames_for_model_bytes:
@@ -576,9 +588,7 @@ def _convert_video_for_preview_if_needed(path: str) -> str:
         return path # Cannot convert, return original
 
     # Check if the video is already likely browser-compatible (MP4 with H.264/H.265)
-    # This check is a heuristic; direct ffprobe is more reliable but adds overhead.
-    # For a preview, extension check is usually sufficient.
-    if path.lower().endswith((".mp4", ".m4v")): # .mov might still need conversion
+    if path.lower().endswith((".mp4", ".m4v")):
         info = _ffprobe_streams(path)
         if info:
             video_streams = [s for s in info.get("streams", []) if s.get("codec_type") == "video"]
@@ -761,7 +771,6 @@ def create_demo():
                 return img_update_clear, video_update_clear, gallery_update_clear, \
                        gr.update(value="", visible=True), main_path_clear, screenshot_paths_clear, raw_media_path_clear
 
-            raw_bytes_for_analysis = b""
             temp_raw_path_for_analysis = ""
             try:
                 # 1. Fetch raw bytes once for both preview and analysis
@@ -789,6 +798,7 @@ def create_demo():
 
                 if not local_playable_path:
                     # If preview failed, cleanup the temp_raw_path_for_analysis as well
+                    # (though atexit/clear_all_files_and_ui would eventually get it)
                     if temp_raw_path_for_analysis in _temp_preview_files_to_delete:
                         _temp_preview_files_to_delete.remove(temp_raw_path_for_analysis)
                     try: os.remove(temp_raw_path_for_analysis)
@@ -811,6 +821,12 @@ def create_demo():
                            gr.update(value=[], visible=False), gr.update(value="Video preview loaded.", visible=True), \
                            local_playable_path, [], temp_raw_path_for_analysis # Update raw_media_path_state
                 else:
+                    # If local_playable_path exists but is not image/video, clean it up
+                    if local_playable_path in _temp_preview_files_to_delete:
+                        _temp_preview_files_to_delete.remove(local_playable_path)
+                    try: os.remove(local_playable_path)
+                    except Exception as e: print(f"Error during cleanup of unplayable temp file {local_playable_path}: {e}")
+
                     return img_update_clear, video_update_clear, gallery_update_clear, \
                            gr.update(value="Preview load failed: unknown playable format.", visible=True), \
                            main_path_clear, screenshot_paths_clear, raw_media_path_clear
@@ -829,7 +845,7 @@ def create_demo():
         url_input.change(
             fn=load_main_preview_and_clear_old,
             inputs=[url_input],
-            outputs=[preview_image, preview_video, screenshot_gallery, preview_status, main_preview_path_state, screenshot_paths_state, raw_media_path_state]
+            outputs=[preview_image, preview_video, screenshot_gallery, preview_status, main_preview_path_state, raw_media_path_state] # Removed screenshot_paths_state from here as it's cleared anyway.
         )
 
         def worker(url: str, prompt: str, key: str, current_main_preview_path: str, raw_media_path: str, progress=gr.Progress()):
@@ -861,7 +877,8 @@ def create_demo():
                     Image.open(BytesIO(raw_bytes_for_analysis)).verify()
                     is_actually_image_for_analysis = True
                 except UnidentifiedImageError:
-                    is_actually_video_for_analysis = True # Not an identifiable image
+                    # If it's not an identifiable image, assume it might be a video
+                    is_actually_video_for_analysis = True
                 except Exception as e:
                     print(f"Warning: PIL error during image verification for raw analysis media ({raw_media_path}): {e}. Falling back to video.")
                     is_actually_video_for_analysis = True # Other PIL error, treat as video
@@ -870,7 +887,7 @@ def create_demo():
 
                 # --- Video Processing Path ---
                 if is_actually_video_for_analysis:
-                    progress(0.25, desc="Running full‑video analysis")
+                    progress(0.25, desc="Running full-video analysis")
                     result_text, generated_screenshot_paths = analyze_video_cohesive(client, raw_media_path, prompt, progress=progress)
 
                 # --- Image Processing Path ---