Update streamlit_app.py

streamlit_app.py

@@ -33,17 +33,14 @@ except Exception:
 
 load_dotenv()
 
-# Logging (minimal)
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("video_ai")
 logger.propagate = False
 
-# App config
 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
 st.session_state.setdefault("videos", "")
 st.session_state.setdefault("loop_video", False)
 st.session_state.setdefault("uploaded_file", None)
@@ -69,7 +66,6 @@ MODEL_OPTIONS = [
     "custom",
 ]
 
-# Utilities
 def sanitize_filename(path_str: str):
     name = Path(path_str).name
     return name.lower().translate(str.maketrans("", "", string.punctuation)).replace(" ", "_")
@@ -145,7 +141,6 @@ def configure_genai_if_needed():
         logger.exception("Failed to configure genai")
     return True
 
-# Upload & processing
 def upload_video_sdk(filepath: str):
     key = get_effective_api_key()
     if not key:
@@ -201,7 +196,6 @@ def wait_for_processed(file_obj, timeout: int = None, progress_callback=None):
         time.sleep(backoff)
         backoff = min(backoff * 2, 8.0)
 
-# Normalize responses into text
 def _normalize_genai_response(response):
     if response is None:
         return ""
@@ -266,42 +260,6 @@ def _normalize_genai_response(response):
             seen.add(t)
     return "\n\n".join(filtered).strip()
 
-# REST fallback to GenAI Responses API
-def rest_responses_api(prompt_text: str, file_path: str, model: 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")
-    url = "https://generativelanguage.googleapis.com/v1beta2/responses:generate"
-    headers = {"Authorization": f"Bearer {key}"}
-    # Build a simple request that attaches the file as a "file" in multipart/form-data.
-    # Use a minimal JSON payload referencing the file by name in the input.
-    fname = Path(file_path).name
-    input_json = {
-        "model": model,
-        "input": [
-            {
-                "text": prompt_text,
-                "mimeType": mimetypes.guess_type(file_path)[0] or "application/octet-stream",
-                "attachments": [{"contentType": mimetypes.guess_type(file_path)[0] or "application/octet-stream", "name": fname}],
-            }
-        ],
-        "maxOutputTokens": max_tokens,
-    }
-    # Multipart: one part "request" with JSON, another with the file binary.
-    try:
-        with open(file_path, "rb") as f:
-            files = {
-                "request": ("request", json.dumps(input_json), "application/json"),
-                "file": (fname, f, mimetypes.guess_type(file_path)[0] or "application/octet-stream"),
-            }
-            resp = requests.post(url, headers=headers, files=files, timeout=timeout)
-        resp.raise_for_status()
-        data = resp.json()
-        return _normalize_genai_response(data)
-    except Exception as e:
-        raise RuntimeError(f"REST Responses API failed: {e}")
-
-# Generation (supports various SDK shapes + REST fallback)
 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:
@@ -313,7 +271,6 @@ def generate_via_responses_api(prompt_text: str, processed, model_used: str, max
             pass
     fname = file_name_or_id(processed) or None
 
-    # Prepare simple system+user structure
     system_msg = {"role": "system", "content": prompt_text}
     user_msg = {"role": "user", "content": f"Please summarize the attached video: {fname or '[uploaded file]'}."}
 
@@ -327,33 +284,11 @@ def generate_via_responses_api(prompt_text: str, processed, model_used: str, max
         txt = str(e_text).lower()
         return any(k in txt for k in ("internal", "unavailable", "deadlineexceeded", "deadline exceeded", "timeout", "rate limit", "503", "502", "500"))
 
-    # Quick pre-check: if processed is a local path (dictless), prefer REST fallback to ensure attachment works
-    local_file_path = None
-    if isinstance(processed, str) and os.path.exists(processed):
-        local_file_path = processed
-    elif isinstance(processed, dict):
-        # if SDK provided a dict with local path info (rare), try to detect
-        for k in ("path", "name", "filename", "uri"):
-            v = processed.get(k)
-            if isinstance(v, str) and os.path.exists(v):
-                local_file_path = v
-                break
-
     start = time.time()
     last_exc = None
     backoff = 1.0
     attempts = 0
 
-    # If we have a local file path, try REST fallback first for reliable file attachment.
-    if local_file_path:
-        try:
-            if progress_callback:
-                progress_callback("rest-fallback", 0, {"file": local_file_path, "model": model_used})
-            return rest_responses_api(prompt_text, local_file_path, model_used, max_tokens=max_tokens, timeout=timeout, progress_callback=progress_callback)
-        except Exception as e:
-            last_exc = e
-            logger.warning("REST fallback failed; will try SDK: %s", e)
-
     while True:
         for method_name, payload in call_variants:
             attempts += 1
@@ -361,20 +296,16 @@ def generate_via_responses_api(prompt_text: str, processed, model_used: str, max
                 if progress_callback:
                     progress_callback("starting", int(time.time() - start), {"model": model_used, "attempt": attempts, "method": method_name})
 
-                # Preferred modern: genai.responses.generate or genai_responses.generate
                 if genai_responses is not None and hasattr(genai_responses, "generate"):
-                    # Remove None entries from payload
                     payload = {k: v for k, v in payload.items() if v is not None}
                     resp = genai_responses.generate(**payload)
                     text = _normalize_genai_response(resp)
                     if progress_callback:
                         progress_callback("done", int(time.time() - start), {"method": method_name})
-                    # If the model returns a request-for-file style message, try REST fallback
                     if text and ("please provide the video" in text.lower() or "upload the video" in text.lower()):
                         raise RuntimeError("Model indicates it didn't receive the file")
                     return text
 
-                # Older path: genai.Responses.create
                 if hasattr(genai, "Responses") and hasattr(genai.Responses, "create"):
                     payload = {k: v for k, v in payload.items() if v is not None}
                     resp = genai.Responses.create(**payload)  # type: ignore
@@ -385,7 +316,6 @@ def generate_via_responses_api(prompt_text: str, processed, model_used: str, max
                         raise RuntimeError("Model indicates it didn't receive the file")
                     return text
 
-                # Fallback: GenerativeModel API (ChatSession). This SDK's ChatSession.send_message may not accept timeout kw.
                 if hasattr(genai, "GenerativeModel"):
                     try:
                         model_obj = genai.GenerativeModel(model_name=model_used)
@@ -395,9 +325,9 @@ def generate_via_responses_api(prompt_text: str, processed, model_used: str, max
                             if send is None:
                                 raise RuntimeError("ChatSession has no send_message")
                             try:
-                                resp = send(prompt_text, timeout=timeout)  # try with timeout
+                                resp = send(prompt_text, timeout=timeout)
                             except TypeError:
-                                resp = send(prompt_text)  # fallback without timeout
+                                resp = send(prompt_text)
                             text = getattr(resp, "text", None) or str(resp)
                             text = text if text else _normalize_genai_response(resp)
                             if progress_callback:
@@ -415,10 +345,9 @@ def generate_via_responses_api(prompt_text: str, processed, model_used: str, max
                 logger.warning("Generation error (model=%s attempt=%s method=%s): %s", model_used, attempts, method_name, msg)
                 if not is_transient_error(msg):
                     if "No supported response generation method" in msg or "has no attribute" in msg or "didn't receive the file" in msg:
-                        # If it's a file-attachment issue or incompatible SDK, offer a clear upgrade message (but don't spam UI)
                         raise RuntimeError(
                             "Installed google-generativeai package may not expose a compatible Responses API or the SDK didn't attach the file correctly. "
-                            "Try upgrading the SDK: pip install --upgrade google-generativeai, or use the app's REST fallback."
+                            "Try upgrading the SDK: pip install --upgrade google-generativeai."
                         ) from e
                     raise
                 if time.time() - start > timeout:
@@ -426,7 +355,6 @@ def generate_via_responses_api(prompt_text: str, processed, model_used: str, max
                 time.sleep(backoff)
                 backoff = min(backoff * 2, 8.0)
 
-# Trim prompt echoes
 def remove_prompt_echo(prompt: str, text: str, check_len: int = 600, ratio_threshold: float = 0.68):
     if not prompt or not text:
         return text
@@ -446,7 +374,7 @@ def remove_prompt_echo(prompt: str, text: str, check_len: int = 600, ratio_thres
             return b_full[len(ph):].lstrip(" \n:-")
     return text
 
-# UI: reset per new URL value
+# UI reset on URL change
 current_url = st.session_state.get("url", "")
 if current_url != st.session_state.get("last_url_value"):
     st.session_state["videos"] = ""
@@ -464,56 +392,63 @@ if current_url != st.session_state.get("last_url_value"):
     st.session_state["last_url_value"] = current_url
 
 st.sidebar.header("Video Input")
-st.sidebar.text_input("Video URL", key="url", placeholder="https://")
+st.sidebar.text_input("Video URL", key="url_input", placeholder="https://", value=st.session_state.get("url", ""))
 
 settings_exp = st.sidebar.expander("Settings", expanded=False)
-chosen = settings_exp.selectbox("Gemini model", MODEL_OPTIONS, index=MODEL_OPTIONS.index(st.session_state.get("preferred_model", "gemini-2.0-flash-lite")))
+chosen = settings_exp.selectbox("Gemini model", MODEL_OPTIONS, index=MODEL_OPTIONS.index(st.session_state.get("preferred_model", "gemini-2.0-flash-lite")), key="model_select")
 custom_model = ""
-if chosen == "custom":
-    custom_model = settings_exp.text_input("Custom model name", value=st.session_state.get("preferred_model", "gemini-2.0-flash-lite"))
-model_input_value = (custom_model.strip() if chosen == "custom" else chosen).strip()
+if settings_exp.session_state.get("model_select") == "custom":
+    custom_model = settings_exp.text_input("Custom model name", value=st.session_state.get("preferred_model", "gemini-2.0-flash-lite"), key="custom_model")
+model_input_value = (custom_model.strip() if custom_model else settings_exp.session_state.get("model_select")).strip()
 
-settings_exp.text_input("Google API Key", key="api_key", value=os.getenv("GOOGLE_API_KEY", ""), type="password")
+settings_exp.text_input("Google API Key", key="api_key_input", value=st.session_state.get("api_key", ""), type="password")
+st.session_state["api_key"] = settings_exp.session_state.get("api_key_input", st.session_state.get("api_key", ""))
 
 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."
 )
-analysis_prompt = settings_exp.text_area("Enter analysis prompt", value=default_prompt, height=140)
-settings_exp.text_input("Video Password (if needed)", key="video-password", placeholder="password", type="password")
+analysis_prompt = settings_exp.text_area("Enter analysis prompt", value=st.session_state.get("analysis_prompt", default_prompt), height=140, key="analysis_prompt")
+st.session_state["analysis_prompt"] = settings_exp.session_state.get("analysis_prompt", default_prompt)
+
+settings_exp.text_input("Video Password (if needed)", key="video_password_input", placeholder="password", type="password")
 
 settings_exp.number_input(
     "Processing timeout (s)", min_value=60, max_value=3600,
     value=st.session_state.get("processing_timeout", 900), step=30,
-    key="processing_timeout",
+    key="processing_timeout_input",
 )
+st.session_state["processing_timeout"] = settings_exp.session_state.get("processing_timeout_input", st.session_state.get("processing_timeout", 900))
+
 settings_exp.number_input(
     "Generation timeout (s)", min_value=30, max_value=1800,
     value=st.session_state.get("generation_timeout", 300), step=10,
-    key="generation_timeout",
+    key="generation_timeout_input",
 )
+st.session_state["generation_timeout"] = settings_exp.session_state.get("generation_timeout_input", st.session_state.get("generation_timeout", 300))
 
 settings_exp.number_input(
     "Compression threshold (MB)", min_value=10, max_value=2000,
     value=st.session_state.get("compression_threshold_mb", 200), step=10,
-    key="compression_threshold_mb",
+    key="compression_threshold_input",
 )
-settings_exp.caption("Files ≤ threshold are uploaded unchanged. Files > threshold are compressed before upload (tunable).")
+st.session_state["compression_threshold_mb"] = settings_exp.session_state.get("compression_threshold_input", st.session_state.get("compression_threshold_mb", 200))
 
 key_source = "session" if st.session_state.get("api_key") else ".env" if os.getenv("GOOGLE_API_KEY") else "none"
-settings_exp.caption(f"Using API key from: **{key_source}**")
+settings_exp.caption(f"Using API key from: {key_source}")
+
 if not get_effective_api_key():
     settings_exp.warning("No Google API key provided; upload/generation disabled.", icon="⚠️")
 
 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()), key="gen_button")
 with col2:
     pass
 
-if st.sidebar.button("Load Video", use_container_width=True):
+if st.sidebar.button("Load Video", use_container_width=True, key="load_video_btn"):
     try:
-        vpw = st.session_state.get("video-password", "")
-        path = download_video_ytdlp(st.session_state.get("url", ""), str(DATA_DIR), vpw)
+        vpw = settings_exp.session_state.get("video_password_input", "")
+        path = download_video_ytdlp(st.session_state.get("url", settings_exp.session_state.get("url_input", "")), str(DATA_DIR), vpw)
         st.session_state["videos"] = path
         st.session_state["last_loaded_path"] = path
         st.session_state.pop("uploaded_file", None)
@@ -532,10 +467,10 @@ if st.session_state["videos"]:
         st.sidebar.write("Couldn't preview video")
 
     with st.sidebar.expander("Options", expanded=False):
-        loop_checkbox = st.checkbox("Enable Loop", value=st.session_state.get("loop_video", False))
-        st.session_state["loop_video"] = loop_checkbox
+        loop_checkbox = st.checkbox("Enable Loop", value=st.session_state.get("loop_video", False), key="loop_checkbox")
+        st.session_state["loop_video"] = settings_exp.session_state.get("loop_checkbox", st.session_state.get("loop_video", False))
 
-        if st.button("Clear Video(s)"):
+        if st.button("Clear Video(s)", key="clear_videos_btn"):
             st.session_state["videos"] = ""
             st.session_state["last_loaded_path"] = ""
             st.session_state["uploaded_file"] = None
@@ -551,7 +486,7 @@ if st.session_state["videos"]:
 
         try:
             with open(st.session_state["videos"], "rb") as vf:
-                st.download_button("Download Video", data=vf, file_name=sanitize_filename(st.session_state["videos"]), mime="video/mp4", use_container_width=True)
+                st.download_button("Download Video", data=vf, file_name=sanitize_filename(st.session_state["videos"]), mime="video/mp4", use_container_width=True, key="download_video_btn")
         except Exception:
             st.sidebar.error("Failed to prepare download")
 
@@ -599,8 +534,8 @@ if generate_now and not st.session_state.get("busy"):
                     reupload_needed = False
 
                 if reupload_needed:
-                    if not HAS_GENAI and not get_effective_api_key():
-                        raise RuntimeError("google.generativeai SDK not available and no API key; cannot upload")
+                    if not HAS_GENAI or upload_file is None:
+                        raise RuntimeError("google.generativeai SDK or upload support unavailable; cannot upload video. Use SDK with upload_file support.")
                     local_path = current_path
 
                     try:
@@ -620,12 +555,7 @@ if generate_now and not st.session_state.get("busy"):
 
                     with st.spinner(f"Uploading video{' (compressed)' if compressed else ''}..."):
                         try:
-                            # Prefer SDK upload if available, else keep local path for REST fallback
-                            if HAS_GENAI and upload_file is not None:
-                                uploaded = upload_video_sdk(upload_path)
-                            else:
-                                # No SDK upload; retain local path (REST fallback will attach file directly)
-                                uploaded = upload_path
+                            uploaded = upload_video_sdk(upload_path)
                         except Exception as e:
                             st.session_state["last_error"] = f"Upload failed: {e}\n\nTraceback:\n{traceback.format_exc()}"
                             st.error("Upload failed. See Last Error for details.")
@@ -653,7 +583,7 @@ if generate_now and not st.session_state.get("busy"):
                     st.session_state["last_loaded_path"] = current_path
                     st.session_state["file_hash"] = current_hash
 
-                prompt_text = (analysis_prompt.strip() or default_prompt).strip()
+                prompt_text = (st.session_state.get("analysis_prompt", "") or default_prompt).strip()
                 out = ""
                 model_used = model_id
                 max_tokens = 2048 if "2.5" in model_used else 1024
@@ -702,14 +632,12 @@ if generate_now and not st.session_state.get("busy"):
             finally:
                 st.session_state["busy"] = False
 
-# Display existing analysis
 if st.session_state.get("analysis_out"):
     just_loaded_same = (st.session_state.get("last_loaded_path") == st.session_state.get("videos"))
     if not just_loaded_same:
         st.subheader("Analysis Result")
         st.markdown(st.session_state.get("analysis_out"))
 
-# Last error expander
 if st.session_state.get("last_error"):
     with st.expander("Last Error", expanded=False):
         st.write(st.session_state.get("last_error"))