Update src/streamlit_app.py

src/streamlit_app.py

@@ -6,7 +6,7 @@ import base64
 import io
 import threading
 import traceback
-import atexit # Correctly imported
+import atexit
 import time
 import logging
 from dotenv import load_dotenv
@@ -30,6 +30,7 @@ from streamlit_webrtc import (
 # --- Configuration ---
 load_dotenv()
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(threadName)s - %(levelname)s - %(message)s')
+logging.info("Application starting up...")
 
 # Audio configuration
 PYAUDIO_FORMAT = pyaudio.paInt16
@@ -46,7 +47,9 @@ VIDEO_FPS_TO_GEMINI = 2
 VIDEO_API_RESIZE = (1024, 1024)
 
 # !!! IMPORTANT: Verify this model name is correct for the Live API !!!
+# This is a common point of failure for ConnectionClosedError.
 MODEL_NAME = "models/gemini-2.0-flash-live-001"
+logging.info(f"Using Gemini Model: {MODEL_NAME}")
 
 MEDICAL_ASSISTANT_SYSTEM_PROMPT = """You are an AI Medical Assistant. Your primary function is to analyze visual information from the user's camera or screen and respond via voice.
 
@@ -65,12 +68,17 @@ Example of a disclaimer you might use: "As an AI assistant, I can describe what
 """
 
 # --- PyAudio Global Instance and Cleanup ---
-pya = pyaudio.PyAudio()
-def cleanup_pyaudio():
-    logging.info("Terminating PyAudio instance.")
-    if pya:
-        pya.terminate()
-atexit.register(cleanup_pyaudio)
+pya = None
+try:
+    pya = pyaudio.PyAudio()
+    def cleanup_pyaudio():
+        logging.info("Terminating PyAudio instance.")
+        if pya: pya.terminate()
+    atexit.register(cleanup_pyaudio)
+    logging.info("PyAudio initialized successfully.")
+except Exception as e_pyaudio:
+    logging.warning(f"PyAudio initialization failed (expected in some server environments): {e_pyaudio}")
+    pya = None
 
 # --- Global Queues - Declare as None, initialize later ---
 video_frames_to_gemini_q: asyncio.Queue = None
@@ -83,6 +91,7 @@ client = None
 if GEMINI_API_KEY:
     try:
         client = genai.Client(http_options={"api_version": "v1beta"}, api_key=GEMINI_API_KEY)
+        logging.info("Gemini client initialized successfully.")
     except Exception as e:
         st.error(f"Failed to initialize Gemini client: {e}")
         logging.critical(f"Gemini client initialization failed: {e}", exc_info=True)
@@ -92,16 +101,14 @@ else:
     logging.critical("GEMINI_API_KEY not found.")
     st.stop()
 
-# Gemini LiveConnectConfig - Simplified for debugging ConnectionClosedError
-# Removed speech_config temporarily. If connection works now, the issue was voice config.
-# If it still fails, check MODEL_NAME and API Key permissions.
+# Gemini LiveConnectConfig - HIGHLY SIMPLIFIED FOR DEBUGGING ConnectionClosedError
+# Start with the absolute minimum. If this connects, incrementally add back features.
+# If this still fails, the issue is likely MODEL_NAME or API Key/Project permissions.
 LIVE_CONNECT_CONFIG = types.LiveConnectConfig(
-    response_modalities=["audio", "text"],
-    # speech_config=types.SpeechConfig( # Temporarily commented out for debugging
-    #     voice_config=types.VoiceConfig(prebuilt_voice_config=types.PrebuiltVoiceConfig(voice_name="Zephyr"))
-    # ),
+    response_modalities=["text"], # Start with text only
+    # speech_config=None, # Explicitly None or omit
 )
-logging.info(f"Using LiveConnectConfig: {LIVE_CONNECT_CONFIG}") # Log the config being used
+logging.info(f"Attempting connection with highly simplified LiveConnectConfig: {LIVE_CONNECT_CONFIG}")
 
 # --- Backend Gemini Interaction Loop ---
 class GeminiInteractionLoop:
@@ -128,11 +135,13 @@ class GeminiInteractionLoop:
                 await asyncio.sleep(0.1); return None
             try:
                 video_frame = await asyncio.wait_for(video_frames_to_gemini_q.get(), timeout=0.02)
+                if video_frame is None: return None # Sentinel received
                 video_frames_to_gemini_q.task_done(); return video_frame
             except asyncio.TimeoutError: pass
             except Exception as e: logging.error(f"Error getting video from queue: {e}", exc_info=True)
             try:
                 audio_chunk = await asyncio.wait_for(audio_chunks_to_gemini_q.get(), timeout=0.02)
+                if audio_chunk is None: return None # Sentinel received
                 audio_chunks_to_gemini_q.task_done(); return audio_chunk
             except asyncio.TimeoutError: return None
             except Exception as e: logging.error(f"Error getting audio from queue: {e}", exc_info=True); return None
@@ -140,10 +149,11 @@ class GeminiInteractionLoop:
             while self.is_running:
                 if not self.gemini_session: await asyncio.sleep(0.1); continue
                 media_data = await get_media_from_queues()
+                if media_data is None and not self.is_running: break # Sentinel and stop signal
                 if media_data and self.gemini_session and self.is_running:
                     try: await self.gemini_session.send(input=media_data)
                     except Exception as e: logging.error(f"Error sending media chunk to Gemini: {e}", exc_info=True)
-                elif not media_data: await asyncio.sleep(0.05)
+                elif not media_data: await asyncio.sleep(0.05) # No data, yield
         except asyncio.CancelledError: logging.info("Task cancelled: stream_media_to_gemini.")
         finally: logging.info("Task finished: stream_media_to_gemini.")
 
@@ -163,8 +173,8 @@ class GeminiInteractionLoop:
                         if text_response := chunk.text:
                             logging.info(f"Gemini text response: {text_response[:100]}")
                             if 'chat_messages' not in st.session_state: st.session_state.chat_messages = []
-                            # This direct update might cause issues if multiple updates happen before rerun
                             st.session_state.chat_messages = st.session_state.chat_messages + [{"role": "assistant", "content": text_response}]
+                            # Consider using st.rerun() via a thread-safe mechanism if immediate UI update is critical
                 except types.generation_types.StopCandidateException: logging.info("Gemini response stream ended normally.")
                 except Exception as e:
                     if self.is_running: logging.error(f"Error receiving from Gemini: {e}", exc_info=True)
@@ -174,8 +184,10 @@ class GeminiInteractionLoop:
 
     async def play_gemini_audio(self):
         logging.info("Task started: Play Gemini audio responses.")
-        # Note: This task might fail in environments without proper ALSA setup (like HF Spaces default)
-        # Error handling below will log the error but allow other tasks to continue.
+        if pya is None:
+             logging.warning("PyAudio not available. Audio playback task will not run.")
+             return
+
         try:
             while audio_from_gemini_playback_q is None and self.is_running: await asyncio.sleep(0.1)
             if not self.is_running: return
@@ -187,32 +199,33 @@ class GeminiInteractionLoop:
             while self.is_running:
                 try:
                     audio_chunk = await asyncio.wait_for(audio_from_gemini_playback_q.get(), timeout=1.0)
+                    if audio_chunk is None and not self.is_running: break # Sentinel and stop signal
                     if audio_chunk: await asyncio.to_thread(self.playback_stream.write, audio_chunk)
                     if audio_chunk: audio_from_gemini_playback_q.task_done()
                 except asyncio.TimeoutError: continue
                 except Exception as e: logging.error(f"Error playing audio chunk: {e}", exc_info=True); await asyncio.sleep(0.01)
         except Exception as e:
-            logging.error(f"Failed to open or use PyAudio playback stream: {e}", exc_info=True)
-            # Don't set self.is_running = False here, just log the playback failure
+            logging.error(f"Failed to open or use PyAudio playback stream (might be expected in this environment): {e}", exc_info=True)
         finally:
             if self.playback_stream:
                 logging.info("Stopping and closing PyAudio playback stream.")
                 try:
                     await asyncio.to_thread(self.playback_stream.stop_stream)
                     await asyncio.to_thread(self.playback_stream.close)
-                except Exception as e_close:
-                     logging.error(f"Error closing playback stream: {e_close}", exc_info=True)
+                except Exception as e_close: logging.error(f"Error closing playback stream: {e_close}", exc_info=True)
                 self.playback_stream = None
             logging.info("Task finished: play_gemini_audio.")
 
     def signal_stop(self):
         logging.info("Signal to stop GeminiInteractionLoop received.")
         self.is_running = False
-        for q in [video_frames_to_gemini_q, audio_chunks_to_gemini_q, audio_from_gemini_playback_q]:
-            if q:
-                try: q.put_nowait(None)
-                except asyncio.QueueFull: logging.warning(f"Queue was full when trying to put sentinel for stop signal.")
-                except Exception as e: logging.error(f"Error putting sentinel in queue: {e}", exc_info=True)
+        for q_name, q_obj_ref in [("video_q", video_frames_to_gemini_q), 
+                                  ("audio_in_q", audio_chunks_to_gemini_q), 
+                                  ("audio_out_q", audio_from_gemini_playback_q)]:
+            if q_obj_ref:
+                try: q_obj_ref.put_nowait(None)
+                except asyncio.QueueFull: logging.warning(f"Queue {q_name} was full when trying to put sentinel for stop signal.")
+                except Exception as e: logging.error(f"Error putting sentinel in {q_name}: {e}", exc_info=True)
 
     async def run_main_loop(self):
         global video_frames_to_gemini_q, audio_chunks_to_gemini_q, audio_from_gemini_playback_q
@@ -228,11 +241,10 @@ class GeminiInteractionLoop:
 
         if client is None: logging.critical("Gemini client is None in run_main_loop. Aborting."); return
 
-        session_task_group = None # Define before try block
         try:
             async with client.aio.live.connect(model=MODEL_NAME, config=LIVE_CONNECT_CONFIG) as session:
                 self.gemini_session = session
-                logging.info("Gemini session established with API.")
+                logging.info(f"Gemini session established with API for model {MODEL_NAME}.")
                 try:
                     logging.info("Sending system prompt to Gemini...")
                     await self.gemini_session.send(input=MEDICAL_ASSISTANT_SYSTEM_PROMPT, end_of_turn=False)
@@ -241,37 +253,35 @@ class GeminiInteractionLoop:
                     logging.error(f"Failed to send system prompt: {e}", exc_info=True)
                     self.is_running = False; return
 
-                async with asyncio.TaskGroup() as tg:
-                    session_task_group = tg # Assign task group
+                # Python 3.9 does not have asyncio.TaskGroup, so manage tasks individually
+                tasks = []
+                try:
                     logging.info("Creating async tasks for Gemini interaction...")
-                    tg.create_task(self.stream_media_to_gemini(), name="stream_media_to_gemini")
-                    tg.create_task(self.process_gemini_responses(), name="process_gemini_responses")
-                    tg.create_task(self.play_gemini_audio(), name="play_gemini_audio")
-                    logging.info("All Gemini interaction tasks created in TaskGroup.")
-                logging.info("Gemini TaskGroup finished execution.")
+                    tasks.append(asyncio.create_task(self.stream_media_to_gemini(), name="stream_media_to_gemini"))
+                    tasks.append(asyncio.create_task(self.process_gemini_responses(), name="process_gemini_responses"))
+                    tasks.append(asyncio.create_task(self.play_gemini_audio(), name="play_gemini_audio"))
+                    logging.info("All Gemini interaction tasks created.")
+                    await asyncio.gather(*tasks) # Wait for all tasks to complete
+                except Exception as e_gather: # Catch errors from tasks gathered
+                    logging.error(f"Error during asyncio.gather: {e_gather}", exc_info=True)
+                    for task in tasks:
+                        if not task.done(): task.cancel() # Cancel pending tasks
+                    await asyncio.gather(*tasks, return_exceptions=True) # Wait for cancellations
+                logging.info("Gemini interaction tasks finished or cancelled.")
 
         except asyncio.CancelledError: logging.info("GeminiInteractionLoop.run_main_loop() was cancelled.")
-        # Removed ExceptionGroup handling as it's not available in Python 3.9
-        except Exception as e:
-            # Log general exceptions, including the ConnectionClosedError if it happens here
+        except Exception as e: # General catch-all, including ConnectionClosedError
             logging.error(f"Exception in GeminiInteractionLoop run_main_loop: {type(e).__name__}: {e}", exc_info=True)
         finally:
             logging.info("GeminiInteractionLoop.run_main_loop() finishing...")
-            self.is_running = False # Ensure flag is set
-            # Cancel any potentially lingering tasks if TaskGroup didn't exit cleanly (though it should)
-            if session_task_group and not session_task_group._closed:
-                 logging.warning("TaskGroup did not close cleanly, attempting cancellation.")
-                 try:
-                     session_task_group.cancel() # Attempt cancellation if needed
-                 except Exception as e_cancel:
-                     logging.error(f"Error cancelling task group: {e_cancel}")
-
-            self.gemini_session = None # Session closed by async with
-            # Clear global queues
+            self.is_running = False # Ensure flag is set for all tasks
+            self.signal_stop() # Send sentinels again to be sure
+            self.gemini_session = None
+            # Clear global queues by setting them to None
             video_frames_to_gemini_q = None
             audio_chunks_to_gemini_q = None
             audio_from_gemini_playback_q = None
-            logging.info("GeminiInteractionLoop finished and queues cleared.")
+            logging.info("GeminiInteractionLoop finished and global queues set to None.")
 
 # --- WebRTC Media Processors ---
 class VideoProcessor(VideoProcessorBase):
@@ -312,6 +322,8 @@ class AudioProcessor(AudioProcessorBase):
         if audio_chunks_to_gemini_q is None: return
         for frame in audio_frames:
             audio_data = frame.planes[0].to_bytes()
+            # Note: Ensure this mime_type and the actual audio data format (sample rate, channels, bit depth)
+            # are compatible with what the Gemini Live API expects for PCM audio.
             mime_type = f"audio/L16;rate={frame.sample_rate};channels={frame.layout.channels}"
             api_data = {"data": audio_data, "mime_type": mime_type}
             try:
@@ -348,7 +360,6 @@ def run_streamlit_app():
     st.markdown("Utilizing Gemini Live API via WebRTC on Hugging Face Spaces")
     st.info("Remember: This AI cannot provide medical diagnoses. Always consult a healthcare professional for medical advice.")
 
-
     with st.sidebar:
         st.header("Session Control")
         if not st.session_state.gemini_session_active:
@@ -356,8 +367,6 @@ def run_streamlit_app():
                 st.session_state.gemini_session_active = True
                 st.session_state.chat_messages = [{"role": "system", "content": "Assistant activating. Please allow camera/microphone access in your browser if prompted."}]
                 
-                # Queues will be initialized inside GeminiInteractionLoop's thread
-                
                 gemini_loop = GeminiInteractionLoop()
                 st.session_state.gemini_loop_instance = gemini_loop
                 threading.Thread(target=lambda: asyncio.run(gemini_loop.run_main_loop()), name="GeminiLoopThread", daemon=True).start()
@@ -401,7 +410,7 @@ def run_streamlit_app():
 
         if webrtc_ctx.state.playing:
             st.caption("WebRTC connected. Streaming your camera and microphone.")
-        elif st.session_state.gemini_session_active:
+        elif st.session_state.gemini_session_active: # Check if session is supposed to be active
             st.caption("WebRTC attempting to connect. Ensure camera/microphone permissions are granted in your browser.")
             if hasattr(webrtc_ctx.state, 'error') and webrtc_ctx.state.error:
                 st.error(f"WebRTC Connection Error: {webrtc_ctx.state.error}")
@@ -409,16 +418,13 @@ def run_streamlit_app():
         st.info("Click 'Start Session' in the sidebar to enable the live feed and assistant.")
 
     st.subheader("Chat with Assistant")
-    # Use a container with a fixed height for scrollable chat
     chat_container = st.container()
     with chat_container:
-        # Display messages from session state
         messages = st.session_state.get('chat_messages', [])
         for msg in messages:
             with st.chat_message(msg["role"]):
                 st.write(msg["content"])
-
-    # Chat input outside the container
+    
     user_chat_input = st.chat_input(
         "Type your message...",
         key="user_chat_input_box",
@@ -426,39 +432,28 @@ def run_streamlit_app():
     )
 
     if user_chat_input:
-        # Append user message immediately for responsiveness
-        st.session_state.chat_messages = st.session_state.get('chat_messages', []) + [{"role": "user", "content": user_chat_input}]
+        current_messages = st.session_state.get('chat_messages', [])
+        current_messages.append({"role": "user", "content": user_chat_input})
+        st.session_state.chat_messages = current_messages
         
-        # Send to backend
         loop_instance = st.session_state.get('gemini_loop_instance')
         if loop_instance and loop_instance.async_event_loop and loop_instance.gemini_session:
             if loop_instance.async_event_loop.is_running():
-                # Use run_coroutine_threadsafe to call async func from Streamlit thread
                 future = asyncio.run_coroutine_threadsafe(
                     loop_instance.send_text_input_to_gemini(user_chat_input),
                     loop_instance.async_event_loop
                 )
-                try:
-                    future.result(timeout=2) # Optional: wait briefly for confirmation
-                except TimeoutError:
-                    logging.warning("Timed out waiting for send_text_input_to_gemini confirmation.")
-                except Exception as e:
-                     logging.error(f"Error calling send_text_input_to_gemini: {e}", exc_info=True)
-
+                try: future.result(timeout=2) 
+                except TimeoutError: logging.warning("Timed out waiting for send_text_input_to_gemini confirmation.")
+                except Exception as e: logging.error(f"Error calling send_text_input_to_gemini: {e}", exc_info=True)
             else: st.error("Session event loop is not running. Cannot send message.")
         elif not loop_instance or not st.session_state.gemini_session_active:
             st.error("Session is not active. Please start a session to send messages.")
         else: st.warning("Session components not fully ready. Please wait a moment.")
-        
-        # Rerun to display the user message and potentially any quick text response from AI
         st.rerun()
 
 if __name__ == "__main__":
-    # Final check before running the app
     if client is None:
-        # Log critical error if client is still None (should have been caught earlier)
         logging.critical("Gemini client could not be initialized. Application cannot start.")
-        # Display error in Streamlit if possible (might not render if client init failed early)
-        st.error("CRITICAL ERROR: Gemini client failed to initialize. Check API Key and logs.")
     else:
         run_streamlit_app()