Update streamlit_app.py

streamlit_app.py

@@ -36,6 +36,7 @@ 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)
@@ -46,6 +47,9 @@ st.session_state.setdefault("analysis_out", "")
 st.session_state.setdefault("last_error", "")
 st.session_state.setdefault("file_hash", None)
 st.session_state.setdefault("fast_mode", False)
+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})
 
 def sanitize_filename(path_str: str):
     name = Path(path_str).name
@@ -101,48 +105,36 @@ def file_name_or_id(file_obj):
         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)
 
-if os.getenv("GOOGLE_API_KEY") and HAS_GENAI:
+def get_effective_api_key():
+    return st.session_state.get("api_key") or os.getenv("GOOGLE_API_KEY")
+
+def configure_genai_if_needed():
+    key = get_effective_api_key()
+    if not key:
+        return False
     try:
-        genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
+        genai.configure(api_key=key)
     except Exception:
+        # ignore here; callers will handle failures
         pass
-
-st.sidebar.header("Video Input")
-st.sidebar.text_input("Video URL", key="url", placeholder="https://")
-
-settings_exp = st.sidebar.expander("Settings", expanded=False)
-env_api_key = os.getenv("GOOGLE_API_KEY", "")
-API_KEY = settings_exp.text_input("Google API Key", value=env_api_key, placeholder="Set GOOGLE_API_KEY in .env or enter here", type="password")
-model_input = settings_exp.text_input("Gemini Model (short name)", "gemini-2.0-flash-lite")
-model_id = model_input.strip() or "gemini-2.0-flash-lite"
-model_arg = model_id if not model_id.startswith("models/") else model_id.split("/", 1)[1]
-default_prompt = (
-    "Watch the video and provide a detailed behavioral report focusing on human actions, interactions, posture, movement, and apparent intent. "
-    "Use vivid, anatomically rich descriptions with numeric estimates for measurements. Include a list of detailed anatomical observations and measurements. "
-    "Adopt a playful, inquisitive persona and ensure the report is engaging and informative."
-)
-analysis_prompt = settings_exp.text_area("Enter analysis", value=default_prompt, height=140)
-settings_exp.text_input("Video Password (if needed)", key="video-password", placeholder="password", type="password")
-settings_exp.checkbox("Fast mode (skip compression, smaller model, fewer tokens)", key="fast_mode")
-
-if not API_KEY and not os.getenv("GOOGLE_API_KEY"):
-    settings_exp.warning("No Google API key provided; upload/generation disabled.", icon="⚠️")
-
-safety_settings = [
-    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"},
-    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"},
-    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "OFF"},
-    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "OFF"},
-]
+    return True
 
 _agent = None
-if HAS_PHI and HAS_GENAI and (API_KEY or os.getenv("GOOGLE_API_KEY")):
+def maybe_create_agent(model_id: str):
+    global _agent
+    key = get_effective_api_key()
+    if not (HAS_PHI and HAS_GENAI and key):
+        _agent = None
+        return None
+    if _agent and st.session_state.get("last_model") == model_id:
+        return _agent
     try:
-        key_to_use = API_KEY if API_KEY else os.getenv("GOOGLE_API_KEY")
-        genai.configure(api_key=key_to_use)
+        genai.configure(api_key=key)
         _agent = Agent(name="Video AI summarizer", model=Gemini(id=model_id), tools=[DuckDuckGo()], markdown=True)
+        st.session_state["last_model"] = model_id
     except Exception:
         _agent = None
+    return _agent
 
 def clear_all_video_state():
     st.session_state.pop("uploaded_file", None)
@@ -165,46 +157,43 @@ if current_url != st.session_state.get("last_url_value"):
     clear_all_video_state()
     st.session_state["last_url_value"] = current_url
 
-if st.sidebar.button("Load Video", use_container_width=True):
-    try:
-        vpw = st.session_state.get("video-password", "")
-        path = download_video_ytdlp(st.session_state.get("url", ""), str(DATA_DIR), vpw)
-        st.session_state["videos"] = path
-        st.session_state["last_loaded_path"] = path
-        st.session_state.pop("uploaded_file", None)
-        st.session_state.pop("processed_file", None)
-        st.session_state["file_hash"] = file_sha256(path)
-    except Exception as e:
-        st.sidebar.error(f"Failed to load video: {e}")
-
-if st.session_state["videos"]:
-    try:
-        st.sidebar.video(st.session_state["videos"], loop=st.session_state.get("loop_video", False))
-    except Exception:
-        st.sidebar.write("Couldn't preview video")
+st.sidebar.header("Video Input")
+st.sidebar.text_input("Video URL", key="url", placeholder="https://")
 
-    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
+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")
+# session API key widget (session-first, fallback to .env)
+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 and avoid anatomically explicit or sensitive detail. Include a list of observations and any timestamps for notable events."
+)
+analysis_prompt = settings_exp.text_area("Enter analysis", value=default_prompt, height=140)
+settings_exp.text_input("Video Password (if needed)", key="video-password", placeholder="password", type="password")
+settings_exp.checkbox("Fast mode (skip compression, smaller model, fewer tokens)", key="fast_mode")
 
-        if st.button("Clear Video(s)"):
-            clear_all_video_state()
+# Show which key is active
+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}**")
 
-        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)
-        except Exception:
-            st.sidebar.error("Failed to prepare download")
+if not get_effective_api_key():
+    settings_exp.warning("No Google API key provided; upload/generation disabled.", icon="⚠️")
 
-    st.sidebar.write("Title:", Path(st.session_state["videos"]).name)
+safety_settings = [
+    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"},
+    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"},
+    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "OFF"},
+    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "OFF"},
+]
 
 def upload_video_sdk(filepath: str):
-    key = API_KEY if API_KEY else os.getenv("GOOGLE_API_KEY")
+    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)
+    # upload_file may stream; wrap to update session progress if supported
     return upload_file(filepath)
 
 def wait_for_processed(file_obj, timeout=180):
@@ -246,20 +235,75 @@ def remove_prompt_echo(prompt: str, text: str, check_len: int = 600, ratio_thres
 
 col1, col2 = st.columns([1, 3])
 with col1:
-    generate_now = st.button("Generate the story", type="primary")
+    generate_now = st.button("Generate the story", type="primary", disabled=not bool(get_effective_api_key()))
 with col2:
     pass
 
+if st.sidebar.button("Load Video", use_container_width=True):
+    try:
+        vpw = st.session_state.get("video-password", "")
+        path = download_video_ytdlp(st.session_state.get("url", ""), str(DATA_DIR), vpw)
+        st.session_state["videos"] = path
+        st.session_state["last_loaded_path"] = path
+        st.session_state.pop("uploaded_file", None)
+        st.session_state.pop("processed_file", None)
+        st.session_state["file_hash"] = file_sha256(path)
+    except Exception as e:
+        st.sidebar.error(f"Failed to load video: {e}")
+
+if st.session_state["videos"]:
+    try:
+        st.sidebar.video(st.session_state["videos"], loop=st.session_state.get("loop_video", False))
+    except Exception:
+        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
+
+        if st.button("Clear Video(s)"):
+            clear_all_video_state()
+
+        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)
+        except Exception:
+            st.sidebar.error("Failed to prepare download")
+
+    st.sidebar.write("Title:", Path(st.session_state["videos"]).name)
+    # show file size and compression suggestion
+    try:
+        file_size_mb = os.path.getsize(st.session_state["videos"]) / (1024 * 1024)
+        st.sidebar.caption(f"File size: {file_size_mb:.1f} MB")
+        if file_size_mb > 50 and not st.session_state.get("fast_mode", False):
+            st.sidebar.warning("Large file detected — consider enabling Fast mode or compression.", icon="⚠️")
+            st.session_state["fast_mode"] = True
+    except Exception:
+        pass
+
 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.")
     else:
-        key_to_use = API_KEY if API_KEY else os.getenv("GOOGLE_API_KEY")
+        key_to_use = get_effective_api_key()
         if not key_to_use:
             st.error("Google API key not set.")
         else:
             try:
                 st.session_state["busy"] = True
+                # ensure genai is configured now
+                try:
+                    genai.configure(api_key=key_to_use)
+                except Exception:
+                    pass
+
+                # recreate/clear agent if key or model changed
+                model_id = (st.session_state.get("model_input") or "gemini-2.0-flash-lite").strip()
+                if st.session_state.get("last_model") != model_id:
+                    # clear cached agent to rebuild with new model/key
+                    st.session_state["last_model"] = ""
+                maybe_create_agent(model_id)
+
                 processed = st.session_state.get("processed_file")
                 current_path = st.session_state.get("videos")
                 try:
@@ -302,15 +346,21 @@ if generate_now and not st.session_state.get("busy"):
 
                 out = ""
                 if st.session_state.get("fast_mode"):
-                    model_used = model_arg if model_arg else "gemini-2.0-flash-lite"
+                    model_used = model_id if model_id else "gemini-2.0-flash-lite"
                     max_tokens = 512
                 else:
-                    model_used = model_arg
+                    model_used = model_id
                     max_tokens = 1024
 
-                if _agent:
+                # cost/tokens estimate (very rough)
+                est_tokens = max_tokens
+                est_cost_caption = f"Est. max tokens: {est_tokens}"
+
+                # Generate via Agent if available
+                agent = maybe_create_agent(model_used)
+                if agent:
                     with st.spinner("Generating description via Agent..."):
-                        response = _agent.run(prompt_text, videos=[processed], safety_settings=safety_settings)
+                        response = agent.run(prompt_text, videos=[processed], safety_settings=safety_settings)
                         out = getattr(response, "content", None) or getattr(response, "outputText", None) or str(response)
                 else:
                     if not HAS_GENAI or genai is None:
@@ -379,6 +429,7 @@ if generate_now and not st.session_state.get("busy"):
                 st.session_state["last_error"] = ""
                 st.subheader("Analysis Result")
                 st.markdown(out)
+                st.caption(est_cost_caption)
             except Exception as e:
                 st.session_state["last_error"] = str(e)
                 st.error("An error occurred while generating the story. You can try Generate again; the uploaded video will be reused.")