Update app.py

app.py

@@ -6,16 +6,17 @@ import time
 import mimetypes
 from pathlib import Path
 
-# Step 1: Enter API Key
-st.sidebar.header("Step 1: Enter Gemini API Key")
-user_api_key = st.sidebar.text_input("Enter your Gemini API key (optional):", type="password")
-
-# Gemini API key configuration
-GEMINI_API_KEY = user_api_key if user_api_key else "AIzaSyDCMPwXHagWqYTQB3HL7FceHEmKUv3v4wc"
-genai.configure(api_key=GEMINI_API_KEY)
-
+# --- Get API key from environment variable or user input ---
+def get_api_key():
+    GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY", "")
+    if not GOOGLE_API_KEY:
+        GOOGLE_API_KEY = st.sidebar.text_input("Enter your Gemini API key", type="password")
+    return GOOGLE_API_KEY or "AIzaSyDCMPwXHagWqYTQB3HL7FceHEmKUv3v4wc"
+
+# VideoProcessor class
 class VideoProcessor:
-    def __init__(self):
+    def __init__(self, api_key):
+        genai.configure(api_key=api_key)
         self.model = genai.GenerativeModel("gemini-2.0-flash")
 
     def upload_video(self, video_path, display_name="uploaded_video"):
@@ -38,64 +39,76 @@ for key in ["video_processor", "video_file", "video_name", "messages"]:
     if key not in st.session_state:
         st.session_state[key] = None if key != "messages" else []
 
-st.title("🎬 Video Retrieval-Augmented Generation - Gemini 2.0")
-
-# Step 2: Upload Video
-st.header("Step 2: Upload Video")
-uploaded_file = st.file_uploader("Choose a video file", type=['mp4', 'mov', 'avi', 'mkv', 'webm'])
-
-if uploaded_file:
-    if mimetypes.guess_type(uploaded_file.name)[0].startswith("video"):
-        file_size = len(uploaded_file.getvalue()) / (1024**2)
-        st.info(f"Size: {file_size:.2f} MB")
-
-        if st.session_state.video_name != uploaded_file.name:
-            st.session_state.video_processor = VideoProcessor()
-            with tempfile.NamedTemporaryFile(delete=False, suffix=Path(uploaded_file.name).suffix) as tmp:
-                tmp.write(uploaded_file.getvalue())
-                tmp_path = tmp.name
-
-            with st.spinner("Uploading and processing..."):
-                video_file = st.session_state.video_processor.upload_video(tmp_path, uploaded_file.name)
-                processed_file = st.session_state.video_processor.wait_for_processing(video_file)
-                st.session_state.video_file = processed_file
-                st.session_state.video_name = uploaded_file.name
-                st.session_state.messages.clear()
-                st.success("✅ Video processed")
-
-            os.unlink(tmp_path)
-
-        st.video(uploaded_file.getvalue())
+# Main app function
+def main():
+    st.set_page_config(page_title="Video Retrieval-Augmented Generation", page_icon="🎬", layout="wide")
+    st.header("Retrieval-Augmented Generation - Gemini 2.0")
+    st.markdown("---")
+
+    api_key = get_api_key()
+
+    if not api_key:
+        st.sidebar.error("Please enter your API key to proceed.")
+        st.stop()
+
+    # Step 2: Upload Video
+    st.subheader("Step 2: Upload your video file")
+    uploaded_file = st.file_uploader("Upload a video", type=['mp4', 'mov', 'avi', 'mkv', 'webm'])
+
+    if uploaded_file:
+        if mimetypes.guess_type(uploaded_file.name)[0].startswith("video"):
+            file_size = len(uploaded_file.getvalue()) / (1024**2)
+            st.info(f"Size: {file_size:.2f} MB")
+
+            if st.session_state.video_name != uploaded_file.name:
+                st.session_state.video_processor = VideoProcessor(api_key)
+                with tempfile.NamedTemporaryFile(delete=False, suffix=Path(uploaded_file.name).suffix) as tmp:
+                    tmp.write(uploaded_file.getvalue())
+                    tmp_path = tmp.name
+
+                with st.spinner("Uploading and processing..."):
+                    video_file = st.session_state.video_processor.upload_video(tmp_path, uploaded_file.name)
+                    processed_file = st.session_state.video_processor.wait_for_processing(video_file)
+                    st.session_state.video_file = processed_file
+                    st.session_state.video_name = uploaded_file.name
+                    st.session_state.messages.clear()
+                    st.success("✅ Video processed")
+
+                os.unlink(tmp_path)
+
+            st.video(uploaded_file.getvalue())
+        else:
+            st.error("Not a valid video file")
+
+    if st.button("Reset Chat"):
+        st.session_state.messages.clear()
+
+    if st.button("Reset All"):
+        st.session_state.clear()
+
+    # Step 3: Chat about Video
+    st.subheader("Step 3: Chat with your video")
+
+    if st.session_state.video_file:
+        for msg in st.session_state.messages:
+            with st.chat_message(msg["role"]):
+                st.markdown(msg["content"])
+
+        user_question = st.chat_input("Ask a question about the video...")
+        if user_question:
+            st.session_state.messages.append({"role": "user", "content": user_question})
+            with st.chat_message("user"):
+                st.markdown(user_question)
+
+            with st.chat_message("assistant"):
+                placeholder = st.empty()
+                with st.spinner("Generating response..."):
+                    response = st.session_state.video_processor.chat_with_video(st.session_state.video_file, user_question)
+
+                placeholder.markdown(response)
+                st.session_state.messages.append({"role": "assistant", "content": response})
     else:
-        st.error("Not a valid video file")
-
-if st.button("Reset Chat"):
-    st.session_state.messages.clear()
-
-if st.button("Reset All"):
-    st.session_state.clear()
-
-# Step 3: Chatting about the Video
-if st.session_state.video_file:
-    st.header("Step 3: Chat with Video")
-
-    for msg in st.session_state.messages:
-        with st.chat_message(msg["role"]):
-            st.markdown(msg["content"])
-
-    prompt = st.chat_input("Ask about the video...")
-
-    if prompt:
-        st.session_state.messages.append({"role": "user", "content": prompt})
-        with st.chat_message("user"):
-            st.markdown(prompt)
-
-        with st.chat_message("assistant"):
-            placeholder = st.empty()
-            with st.spinner("Thinking..."):
-                response = st.session_state.video_processor.chat_with_video(st.session_state.video_file, prompt)
+        st.info("👈 Please upload a video to start chatting.")
 
-            placeholder.markdown(response)
-            st.session_state.messages.append({"role": "assistant", "content": response})
-else:
-    st.info("👈 Please upload a video to start chatting.")
+if __name__ == "__main__":
+    main()