Update streamlit_app.py

streamlit_app.py

@@ -1,4 +1,4 @@
-# streamlit_app.py
+# streamlit_app_enhanced.py
 import os
 import time
 import string
@@ -9,6 +9,7 @@ from pathlib import Path
 from difflib import SequenceMatcher
 import concurrent.futures
 import json
+import logging
 
 import yt_dlp
 import ffmpeg
@@ -37,11 +38,15 @@ except Exception:
     upload_file = get_file = None
     HAS_GENAI = False
 
+# Logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("video_ai")
+
 st.set_page_config(page_title="Generate the story of videos", layout="wide")
 DATA_DIR = Path("./data")
 DATA_DIR.mkdir(exist_ok=True)
 
-# ---- Session defaults ----
+# Session defaults
 st.session_state.setdefault("videos", "")
 st.session_state.setdefault("loop_video", False)
 st.session_state.setdefault("uploaded_file", None)
@@ -55,10 +60,11 @@ st.session_state.setdefault("api_key", os.getenv("GOOGLE_API_KEY", ""))
 st.session_state.setdefault("last_model", "")
 st.session_state.setdefault("upload_progress", {"uploaded": 0, "total": 0})
 st.session_state.setdefault("last_url_value", "")
-st.session_state.setdefault("processing_timeout", 900)  # increased default to 15m
-st.session_state.setdefault("generation_timeout", 300)  # for Responses generate
+st.session_state.setdefault("processing_timeout", 900)
+st.session_state.setdefault("generation_timeout", 300)
+st.session_state.setdefault("preferred_model", "gemini-2.5-flash-lite")
 
-# ---- Helpers ----
+# Helpers (kept in-file for single-file deliverable)
 def sanitize_filename(path_str: str):
     name = Path(path_str).name
     return name.lower().translate(str.maketrans("", "", string.punctuation)).replace(" ", "_")
@@ -129,7 +135,7 @@ def configure_genai_if_needed():
         pass
     return True
 
-# ---- Agent management ----
+# Agent management
 _agent = None
 def maybe_create_agent(model_id: str):
     global _agent
@@ -167,12 +173,12 @@ if current_url != st.session_state.get("last_url_value"):
     clear_all_video_state()
     st.session_state["last_url_value"] = current_url
 
-# ---- Sidebar UI ----
+# Sidebar UI
 st.sidebar.header("Video Input")
 st.sidebar.text_input("Video URL", key="url", placeholder="https://")
 
 settings_exp = st.sidebar.expander("Settings", expanded=False)
-model_input = settings_exp.text_input("Gemini Model (short name)", "gemini-2.0-flash-lite", key="model_input")
+model_input = settings_exp.text_input("Preferred Gemini Model (short name)", st.session_state.get("preferred_model", "gemini-2.5-flash-lite"), key="model_input")
 settings_exp.text_input("Google API Key", key="api_key", value=os.getenv("GOOGLE_API_KEY", ""), type="password")
 default_prompt = (
     "Watch the video and provide a detailed behavioral report focusing on human actions, interactions, posture, movement, and apparent intent. Keep language professional. Include a list of observations for notable events."
@@ -203,21 +209,22 @@ safety_settings = [
     {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "OFF"},
 ]
 
-# ---- Upload & processing helpers ----
-def upload_video_sdk(filepath: str):
+# Upload & processing helpers
+def upload_video_sdk(filepath: str, progress_callback=None):
     key = get_effective_api_key()
     if not key:
         raise RuntimeError("No API key provided")
     if not HAS_GENAI or upload_file is None:
         raise RuntimeError("google.generativeai SDK not available; cannot upload")
     genai.configure(api_key=key)
-    return upload_file(filepath)
+    # upload_file doesn't offer progress hooks in SDK; attempt best-effort by streaming in chunks if possible
+    # Fall back to direct upload_file call for compatibility
+    try:
+        return upload_file(filepath)
+    except Exception as e:
+        raise
 
-def wait_for_processed(file_obj, timeout: int = None):
-    """
-    Poll get_file until file is no longer PROCESSING.
-    Retries get_file on transient errors with exponential backoff.
-    """
+def wait_for_processed(file_obj, timeout: int = None, progress_callback=None):
     if timeout is None:
         timeout = st.session_state.get("processing_timeout", 900)
     if not HAS_GENAI or get_file is None:
@@ -227,6 +234,7 @@ def wait_for_processed(file_obj, timeout: int = None):
     if not name:
         return file_obj
     backoff = 1.0
+    last_state = None
     while True:
         try:
             obj = get_file(name)
@@ -238,13 +246,24 @@ def wait_for_processed(file_obj, timeout: int = None):
             continue
 
         state = getattr(obj, "state", None)
-        if not state or getattr(state, "name", None) != "PROCESSING":
+        state_name = getattr(state, "name", None) if state else None
+        if progress_callback:
+            # show a simple heuristic percent while PROCESSING
+            elapsed = int(time.time() - start)
+            pct = 100 if not state_name else (50 if state_name == "PROCESSING" else 100)
+            try:
+                progress_callback(min(100, pct), elapsed, state_name)
+            except Exception:
+                pass
+
+        if not state_name or state_name != "PROCESSING":
             return obj
 
         if time.time() - start > timeout:
             raise TimeoutError(f"File processing timed out after {int(time.time() - start)}s")
         time.sleep(backoff)
         backoff = min(backoff * 2, 8.0)
+        last_state = state_name
 
 def remove_prompt_echo(prompt: str, text: str, check_len: int = 600, ratio_threshold: float = 0.68):
     if not prompt or not text:
@@ -285,8 +304,8 @@ def compress_video_if_large(local_path: str, threshold_mb: int = 50):
         st.session_state["last_error"] = f"Video compression failed: {e}\n{traceback.format_exc()}"
         return local_path, False
 
-# ---- Robust Responses API caller adapted for varying model versions ----
-def generate_via_responses_api(prompt_text: str, processed, model_used: str, max_tokens: int = 1024, timeout: int = 300):
+# Robust Responses API caller with retries and auto-fallback to older model on certain failures
+def generate_via_responses_api(prompt_text: str, processed, model_used: str, max_tokens: int = 1024, timeout: int = 300, progress_callback=None):
     key = get_effective_api_key()
     if not key:
         raise RuntimeError("No API key provided")
@@ -300,79 +319,98 @@ def generate_via_responses_api(prompt_text: str, processed, model_used: str, max
     system_msg = {"role": "system", "content": prompt_text}
     user_msg = {"role": "user", "content": "Please summarize the attached video."}
 
-    # Some model versions and SDK releases expect messages, some older ones expect input with files.
     call_variants = [
         {"messages": [system_msg, user_msg], "files": [{"name": fname}], "safety_settings": safety_settings, "max_output_tokens": max_tokens},
         {"input": [{"text": prompt_text, "files": [{"name": fname}]}], "safety_settings": safety_settings, "max_output_tokens": max_tokens},
     ]
 
-    last_exc = None
+    def is_transient_error(e_text: str):
+        txt = str(e_text).lower()
+        return any(k in txt for k in ("internal", "unavailable", "deadlineexceeded", "deadline exceeded", "timeout", "rate limit", "503", "502", "500"))
+
     start = time.time()
+    last_exc = None
     backoff = 1.0
-    while True:
-        for payload in call_variants:
-            try:
-                response = genai.responses.generate(model=model_used, **payload)
-                # If successful, normalize below
-                return _normalize_genai_response(response)
-            except Exception as e:
-                last_exc = e
-                # If it's a transient server error, let outer retry/backoff handle it
-                # Quick heuristic: inspect message for INTERNAL/UNAVAILABLE/DeadlineExceeded
-                msg = str(e).lower()
-                if any(k in msg for k in ("internal", "unavailable", "deadlineexceeded", "deadline exceeded", "timeout", "rate limit")):
-                    # will retry below
-                    pass
-                else:
-                    # If it's a clear invalid-argument or permission error, bubble up immediately
-                    raise
-        if time.time() - start > timeout:
-            raise TimeoutError(f"Responses.generate timed out after {timeout}s: last error: {last_exc}")
-        time.sleep(backoff)
-        backoff = min(backoff * 2, 8.0)
+    max_total = timeout
+    attempts = 0
+    tried_models = []
+    preferred_model = model_used or st.session_state.get("preferred_model", "gemini-2.5-flash-lite")
+    fallback_model = "gemini-2.0-flash-lite" if "2.5" in preferred_model else None
+    models_to_try = [preferred_model] + ([fallback_model] if fallback_model else [])
+    for m in models_to_try:
+        if not m:
+            continue
+        tried_models.append(m)
+        # per-model attempt window
+        model_start = time.time()
+        while True:
+            attempts += 1
+            for payload in call_variants:
+                try:
+                    if progress_callback:
+                        elapsed = int(time.time() - start)
+                        try:
+                            progress_callback("starting_generation", elapsed, {"model": m, "attempt": attempts})
+                        except Exception:
+                            pass
+                    response = genai.responses.generate(model=m, **payload)
+                    text = _normalize_genai_response(response)
+                    if progress_callback:
+                        elapsed = int(time.time() - start)
+                        try:
+                            progress_callback("generation_complete", elapsed, {"model": m})
+                        except Exception:
+                            pass
+                    return text
+                except Exception as e:
+                    last_exc = e
+                    msg = str(e)
+                    logger.warning("Responses.generate error on model %s attempt %s: %s", m, attempts, msg)
+                    if not is_transient_error(msg):
+                        # Non-transient: rethrow to surface to caller
+                        raise
+                    # transient: will retry for this model up to timeout
+                    if time.time() - start > max_total:
+                        break
+                    time.sleep(backoff)
+                    backoff = min(backoff * 2, 8.0)
+            if time.time() - model_start > max_total:
+                break
+        # try next model (fallback)
+    raise TimeoutError(f"Responses.generate failed after trying models {tried_models}: last error: {last_exc}")
 
 def _normalize_genai_response(response):
-    # Accept dict or object shapes. Extract text pieces robustly and join.
     outputs = []
     if response is None:
         return ""
-
-    # If it's an object with attributes
     if not isinstance(response, dict):
         try:
             response = json.loads(str(response))
         except Exception:
-            # fallback to attribute access
             pass
-
-    # Strategy: check common keys
     candidate_lists = []
-    for key in ("output", "candidates", "items", "responses", "choices"):
-        val = response.get(key) if isinstance(response, dict) else None
-        if isinstance(val, list) and val:
-            candidate_lists.append(val)
-    if not candidate_lists:
-        # fallback: any list value
-        if isinstance(response, dict):
-            for v in response.values():
-                if isinstance(v, list) and v:
-                    candidate_lists.append(v)
-                    break
-
+    if isinstance(response, dict):
+        for key in ("output", "candidates", "items", "responses", "choices"):
+            val = response.get(key)
+            if isinstance(val, list) and val:
+                candidate_lists.append(val)
+    if not candidate_lists and isinstance(response, dict):
+        for v in response.values():
+            if isinstance(v, list) and v:
+                candidate_lists.append(v)
+                break
     text_pieces = []
     for lst in candidate_lists:
         for item in lst:
             if not item:
                 continue
             if isinstance(item, dict):
-                # common text keys
                 for k in ("content", "text", "message", "output_text", "output"):
                     t = item.get(k)
                     if t:
                         text_pieces.append(str(t).strip())
                         break
                 else:
-                    # nested forms
                     if "content" in item and isinstance(item["content"], list):
                         for part in item["content"]:
                             if isinstance(part, dict):
@@ -384,23 +422,18 @@ def _normalize_genai_response(response):
             elif isinstance(item, str):
                 text_pieces.append(item.strip())
             else:
-                # try attribute access
                 try:
                     t = getattr(item, "text", None) or getattr(item, "content", None)
                     if t:
                         text_pieces.append(str(t).strip())
                 except Exception:
                     pass
-
-    # If still empty, try top-level text fields
     if not text_pieces and isinstance(response, dict):
         for k in ("text", "message", "output_text"):
             v = response.get(k)
             if v:
                 text_pieces.append(str(v).strip())
                 break
-
-    # deduplicate preserving order
     seen = set()
     filtered = []
     for t in text_pieces:
@@ -411,7 +444,7 @@ def _normalize_genai_response(response):
             seen.add(t)
     return "\n\n".join(filtered).strip()
 
-# ---- Layout ----
+# Layout
 col1, col2 = st.columns([1, 3])
 with col1:
     generate_now = st.button("Generate the story", type="primary", disabled=not bool(get_effective_api_key()))
@@ -461,7 +494,7 @@ if st.session_state["videos"]:
     except Exception:
         pass
 
-# ---- Main generation flow ----
+# Main generation flow
 if generate_now and not st.session_state.get("busy"):
     if not st.session_state.get("videos"):
         st.error("No video loaded. Use 'Load Video' in the sidebar.")
@@ -478,7 +511,7 @@ if generate_now and not st.session_state.get("busy"):
                 except Exception:
                     pass
 
-                model_id = (st.session_state.get("model_input") or "gemini-2.0-flash-lite").strip()
+                model_id = (st.session_state.get("model_input") or st.session_state.get("preferred_model") or "gemini-2.5-flash-lite").strip()
                 if st.session_state.get("last_model") != model_id:
                     st.session_state["last_model"] = ""
                 maybe_create_agent(model_id)
@@ -501,9 +534,7 @@ if generate_now and not st.session_state.get("busy"):
                     upload_path, compressed = compress_video_if_large(local_path)
 
                     with st.spinner(f"Uploading video{' (compressed)' if compressed else ''}..."):
-                        # Provide an upload progress bar UI while calling upload_file.
-                        progress_placeholder = st.empty()
-                        progress_bar = None
+                        upload_progress_placeholder = st.empty()
                         try:
                             uploaded = upload_video_sdk(upload_path)
                         except Exception as e:
@@ -512,11 +543,15 @@ if generate_now and not st.session_state.get("busy"):
                             raise
 
                     try:
-                        # Show a more informative processing progress area
                         processing_placeholder = st.empty()
                         processing_bar = processing_placeholder.progress(0)
-                        start_wait = time.time()
-                        processed = wait_for_processed(uploaded, timeout=st.session_state.get("processing_timeout", 900))
+                        def processing_cb(pct, elapsed, state):
+                            try:
+                                processing_bar.progress(min(100, int(pct)))
+                                processing_placeholder.caption(f"State: {state} — elapsed: {elapsed}s")
+                            except Exception:
+                                pass
+                        processed = wait_for_processed(uploaded, timeout=st.session_state.get("processing_timeout", 900), progress_callback=processing_cb)
                         processing_bar.progress(100)
                         processing_placeholder.success("Processing complete")
                     except Exception as e:
@@ -535,7 +570,6 @@ if generate_now and not st.session_state.get("busy"):
                 max_tokens = 2048 if "2.5" in model_used else 1024
                 est_tokens = max_tokens
 
-                # Try Agent first, fallback to Responses API
                 agent = maybe_create_agent(model_used)
                 debug_info = {"agent_attempted": False, "agent_ok": False, "agent_error": None, "agent_response_has_text": False}
                 if agent:
@@ -567,8 +601,16 @@ if generate_now and not st.session_state.get("busy"):
 
                 if not out:
                     try:
-                        with st.spinner("Generating description via Responses API..."):
-                            out = generate_via_responses_api(prompt_text, processed, model_used, max_tokens=max_tokens, timeout=st.session_state.get("generation_timeout", 300))
+                        gen_progress_placeholder = st.empty()
+                        gen_status = gen_progress_placeholder.text("Starting generation...")
+                        start_gen = time.time()
+                        def gen_progress_cb(stage, elapsed, info):
+                            try:
+                                gen_status.text(f"Stage: {stage} — elapsed: {elapsed}s — {info}")
+                            except Exception:
+                                pass
+                        out = generate_via_responses_api(prompt_text, processed, model_used, max_tokens=max_tokens, timeout=st.session_state.get("generation_timeout", 300), progress_callback=gen_progress_cb)
+                        gen_progress_placeholder.text(f"Generation complete in {int(time.time()-start_gen)}s")
                     except Exception as e:
                         tb = traceback.format_exc()
                         st.session_state["last_error"] = f"Responses API error: {e}\n\nDebug: {debug_info}\n\nTraceback:\n{tb}"