Update streamlit_app.py

streamlit_app.py

@@ -28,8 +28,9 @@ except Exception:
 # google.generativeai SDK (guarded)
 try:
     import google.generativeai as genai
-    from google.generativeai import upload_file, get_file
-
+    # older SDK exposed upload_file/get_file at top-level; try to bind them if present
+    upload_file = getattr(genai, "upload_file", None)
+    get_file = getattr(genai, "get_file", None)
     HAS_GENAI = True
 except Exception:
     genai = None
@@ -94,9 +95,7 @@ def convert_video_to_mp4(video_path: str) -> str:
     try:
         ffmpeg.input(video_path).output(target_path).run(overwrite_output=True, quiet=True)
     except Exception:
-        # re-raise so caller can handle
         raise
-    # remove source only if different and successful
     if os.path.exists(target_path) and os.path.getsize(target_path) > 0:
         try:
             if str(Path(video_path).resolve()) != str(Path(target_path).resolve()):
@@ -152,8 +151,21 @@ def file_name_or_id(file_obj):
     if file_obj is None:
         return None
     if isinstance(file_obj, dict):
-        return file_obj.get("name") or file_obj.get("id")
-    return getattr(file_obj, "name", None) or getattr(file_obj, "id", None) or getattr(file_obj, "fileId", None)
+        return file_obj.get("name") or file_obj.get("id") or file_obj.get("fileId")
+    # SDK objects might expose .name, .id, .fileId, or nested fields
+    for attr in ("name", "id", "fileId", "file_id"):
+        v = getattr(file_obj, attr, None)
+        if v:
+            return v
+    # sometimes SDK returns {'file': {...}}
+    try:
+        if hasattr(file_obj, "to_dict"):
+            d = file_obj.to_dict()
+            if isinstance(d, dict):
+                return d.get("name") or d.get("id") or d.get("fileId")
+    except Exception:
+        pass
+    return None
 
 def get_effective_api_key():
     return st.session_state.get("api_key") or os.getenv("GOOGLE_API_KEY")
@@ -203,7 +215,6 @@ def clear_all_video_state():
 # Reset when URL changes (compare against last_url only)
 current_url = st.session_state.get("url", "")
 if current_url != st.session_state.get("last_url"):
-    # avoid clearing on first load when last_url is empty
     if st.session_state.get("last_url"):
         clear_all_video_state()
     st.session_state["last_url"] = current_url
@@ -254,41 +265,120 @@ safety_settings = [
 ]
 
 # ---- Upload & processing helpers ----
+def _normalize_uploaded_obj(obj):
+    """
+    Normalize possible upload_file return shapes:
+    - SDK object with .name or .id
+    - dict with 'name' or 'id'
+    - legacy or ragstore shapes: {'file': {...}} etc.
+    Return a dict-like wrapper with at least 'name' key or original obj.
+    """
+    if obj is None:
+        return None
+    if isinstance(obj, dict):
+        # common keys
+        name = obj.get("name") or obj.get("id") or obj.get("fileId") or (obj.get("file") and obj["file"].get("name"))
+        if name:
+            return {"name": name, "_raw": obj}
+        return obj
+    # SDK object: try attributes and to_dict
+    for attr in ("name", "id", "fileId", "file_id"):
+        v = getattr(obj, attr, None)
+        if v:
+            return {"name": v, "_raw": obj}
+    try:
+        if hasattr(obj, "to_dict"):
+            d = obj.to_dict()
+            if isinstance(d, dict):
+                name = d.get("name") or d.get("id") or d.get("fileId") or (d.get("file") and d["file"].get("name"))
+                if name:
+                    return {"name": name, "_raw": obj}
+    except Exception:
+        pass
+    return {"name": None, "_raw": obj}
+
 def upload_video_sdk(filepath: str):
     key = get_effective_api_key()
     if not key:
         raise RuntimeError("No API key provided")
-    if not HAS_GENAI or upload_file is None:
+    if not HAS_GENAI or genai is None:
         raise RuntimeError("google.generativeai SDK not available; cannot upload")
+    # prefer upload_file if present, otherwise try genai.ragstore or genai.files.upload style
     genai.configure(api_key=key)
-    return upload_file(filepath)
+    last_exc = None
+    # try a few known entrypoints
+    candidate_calls = []
+    if upload_file:
+        candidate_calls.append(("upload_file", upload_file))
+    # sometimes SDK exposed as genai.ragstore.upload_file or genai.files.upload
+    candidate_calls.append(("genai.ragstore.upload_file", getattr(genai, "ragstore", None) and getattr(genai.ragstore, "upload_file", None)))
+    candidate_calls.append(("genai.files.upload", getattr(genai, "files", None) and getattr(genai.files, "upload", None)))
+    candidate_calls.append(("genai.upload", getattr(genai, "upload", None)))
+    # filter none
+    candidate_calls = [(n, fn) for n, fn in candidate_calls if fn]
+    if not candidate_calls:
+        raise RuntimeError("No upload function available in google.generativeai SDK")
+    for name, fn in candidate_calls:
+        try:
+            res = fn(filepath)
+            norm = _normalize_uploaded_obj(res)
+            # ensure we return something that file_name_or_id understands
+            return norm if isinstance(norm, dict) else res
+        except Exception as e:
+            last_exc = e
+            continue
+    raise RuntimeError(f"All upload methods failed. Last error: {last_exc}")
 
 def wait_for_processed(file_obj, timeout: int = None):
     if timeout is None:
         timeout = st.session_state.get("processing_timeout", 900)
-    if not HAS_GENAI or get_file is None:
+    if not HAS_GENAI or genai is None:
         return file_obj
     start = time.time()
     name = file_name_or_id(file_obj)
+    # If no determinable name/id, just return file_obj
     if not name:
         return file_obj
+    genai.configure(api_key=get_effective_api_key())
     backoff = 1.0
-    while True:
-        try:
-            obj = get_file(name)
-        except Exception:
-            if time.time() - start > timeout:
-                raise TimeoutError("Failed to fetch file status before timeout")
-            time.sleep(backoff)
-            backoff = min(backoff * 2, 8.0)
-            continue
-
-        state = getattr(obj, "state", None)
-        if not state or getattr(state, "name", None) != "PROCESSING":
-            return obj
+    last_exc = None
+    # prefer get_file if present, otherwise try genai.ragstore.get or genai.files.get
+    candidate_getters = []
+    if get_file:
+        candidate_getters.append(("get_file", get_file))
+    candidate_getters.append(("genai.ragstore.get", getattr(genai, "ragstore", None) and getattr(genai.ragstore, "get", None)))
+    candidate_getters.append(("genai.files.get", getattr(genai, "files", None) and getattr(genai.files, "get", None)))
+    candidate_getters = [(n, fn) for n, fn in candidate_getters if fn]
+    # If none, return file_obj immediately
+    if not candidate_getters:
+        return file_obj
 
+    while True:
+        for name_label, getter in candidate_getters:
+            try:
+                obj = getter(name)
+                # normalize possible SDK object/dict
+                # check for state attribute or dict field
+                state = None
+                if isinstance(obj, dict):
+                    state = obj.get("state") or (obj.get("status") and {"name": obj.get("status")})
+                else:
+                    state = getattr(obj, "state", None) or getattr(obj, "status", None)
+                if not state:
+                    return obj
+                # state might be dict or object with .name
+                state_name = state.get("name") if isinstance(state, dict) else getattr(state, "name", None)
+                if state_name and state_name.upper() == "PROCESSING":
+                    # still processing; continue polling
+                    last_obj = obj
+                    break
+                return obj
+            except Exception as e:
+                last_exc = e
+                # transient errors: backoff and retry until timeout
+                continue
         if time.time() - start > timeout:
-            raise TimeoutError("File processing timed out")
+            raise TimeoutError(f"File processing timed out. Last error: {last_exc}")
         time.sleep(backoff)
         backoff = min(backoff * 2, 8.0)
 
@@ -341,9 +431,8 @@ def compress_video_if_large(local_path: str, threshold_mb: int = 200):
 # ---- Simple 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()))    
+    generate_now = st.button("Generate the story", type="primary", disabled=not bool(get_effective_api_key()))
 with col2:
-    # small UX note column
     if not st.session_state.get("videos"):
         st.info("Load a video first (sidebar) to enable generation.", icon="ℹ️")
     else:
@@ -355,7 +444,6 @@ if st.sidebar.button("Load Video", use_container_width=True):
         path = download_video_ytdlp(st.session_state.get("url", ""), str(DATA_DIR), vpw)
         st.session_state["videos"] = path
         st.session_state["last_local_path"] = path
-        # keep last_url intact — it tracks the input URL
         st.session_state.pop("uploaded_file", None)
         st.session_state.pop("processed_file", None)
         try:
@@ -367,7 +455,6 @@ if st.sidebar.button("Load Video", use_container_width=True):
 
 if st.session_state["videos"]:
     path = st.session_state["videos"]
-    # ensure mp4 for preview and read bytes for reliable preview
     try:
         mp4_path = convert_video_to_mp4(path)
         with open(mp4_path, "rb") as vf:
@@ -398,7 +485,7 @@ if st.session_state["videos"]:
     except Exception:
         pass
 
-# ---- Generation flow (minimal, robust) ----
+# ---- Generation flow (fixed and robust) ----
 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.")
@@ -430,7 +517,6 @@ if generate_now and not st.session_state.get("busy"):
                 reupload_needed = True
                 uploaded_file = st.session_state.get("uploaded_file")
                 uploaded_name = file_name_or_id(uploaded_file)
-                # Use last_local_path to determine if file changed
                 if processed and st.session_state.get("last_local_path") == current_path and st.session_state.get("file_hash") == current_hash and uploaded_name:
                     reupload_needed = False
 
@@ -446,7 +532,8 @@ if generate_now and not st.session_state.get("busy"):
                         except Exception as e:
                             st.session_state["last_error"] = f"Upload failed for {upload_path}: {e}\n{traceback.format_exc()}"
                             st.error(f"Upload failed: {e}. Check the error log for more details.")
-                            return  # Prevent further processing.
+                            st.session_state["busy"] = False
+                            raise
 
                     try:
                         processing_placeholder = st.empty()
@@ -461,6 +548,7 @@ if generate_now and not st.session_state.get("busy"):
                     except Exception as e:
                         st.session_state["last_error"] = f"Processing failed/wait timeout: {e}"
                         st.error("Video processing failed or timed out. See Last Error.")
+                        st.session_state["busy"] = False
                         raise
 
                     st.session_state["uploaded_file"] = uploaded
@@ -501,7 +589,6 @@ if generate_now and not st.session_state.get("busy"):
                         debug_info["agent_error"] = f"{ae}"
 
                 if not out:
-                    # Use Responses API directly
                     def generate_via_responses_api(prompt_text: str, processed, model_used: str, max_tokens: int = 1024, timeout: int = 300):
                         key = get_effective_api_key()
                         if not key:
@@ -544,7 +631,6 @@ if generate_now and not st.session_state.get("busy"):
                         outputs = []
                         if response is None:
                             return ""
-                        # response may be SDK object or dict-like; coerce to string chunks
                         text_pieces = []
                         try:
                             if isinstance(response, dict):
@@ -567,6 +653,12 @@ if generate_now and not st.session_state.get("busy"):
                                             text_pieces.append(str(v).strip())
                                             break
                             else:
+                                # SDK object: try attributes
+                                for attr in ("output", "candidates", "text", "message", "content"):
+                                    v = getattr(response, attr, None)
+                                    if v:
+                                        text_pieces.append(str(v).strip())
+                            if not text_pieces:
                                 text_pieces.append(str(response))
                         except Exception:
                             text_pieces.append(str(response))
@@ -611,4 +703,4 @@ if generate_now and not st.session_state.get("busy"):
                 st.session_state["last_error"] = f"{e}\n{traceback.format_exc()}"
                 st.error("An error occurred while generating the story. You can try Generate again; the uploaded video will be reused.")
             finally:
-                st.session_state["busy"] = False
+                st.session_state["busy"] = False