Update src/streamlit_app.py

src/streamlit_app.py

@@ -15,6 +15,9 @@ import cv2
 import pyaudio
 import PIL.Image
 
+# Import websockets for explicit exception handling
+import websockets.exceptions
+
 from google import genai
 from google.genai import types
 
@@ -28,12 +31,11 @@ from streamlit_webrtc import (
 load_dotenv()
 
 
-# Audio configuration
-
+# Audio configuration - fix audio format issues
 FORMAT = pyaudio.paInt16
 CHANNELS = 1
-SEND_SAMPLE_RATE = 16000
-RECEIVE_SAMPLE_RATE = 24000
+SEND_SAMPLE_RATE = 16000  # Changed to match mime_type for consistency
+RECEIVE_SAMPLE_RATE = 16000  # Changed from 24000 to 16000 to match send rate
 CHUNK_SIZE = 1024
 
 # Map PyAudio format to a more descriptive name for clarity.
@@ -43,9 +45,9 @@ PYAUDIO_PLAYBACK_CHUNK_SIZE = CHUNK_SIZE
 GEMINI_AUDIO_RECEIVE_SAMPLE_RATE = RECEIVE_SAMPLE_RATE
 
 # Video configuration
-
-VIDEO_FPS_TO_GEMINI = 2
-VIDEO_API_RESIZE = (1024, 1024)
+VIDEO_FPS_TO_GEMINI = 1  # Reduced from 2 to lower bandwidth
+VIDEO_API_RESIZE = (512, 512)  # Reduced from 1024x1024 to lower payload size
+MAX_PAYLOAD_SIZE_BYTES = 60000  # Just under 64KB WebSocket limit
 
 # Queue sizes
 MEDIA_TO_GEMINI_QUEUE_MAXSIZE = 10
@@ -96,30 +98,54 @@ audio_chunks_to_gemini_q: asyncio.Queue = None
 audio_from_gemini_playback_q: asyncio.Queue = None
 
 # --- Gemini Client Setup ---
-GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
-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)
-        st.stop()
-else:
-    st.error(
-        "GEMINI_API_KEY not found in environment variables. Please set it for the application to run.")
-    logging.critical("GEMINI_API_KEY not found.")
-    st.stop()
+# Try to get API key from environment or use a manually provided one
+def initialize_gemini_client():
+    # Check for API key in various places
+    api_key = os.environ.get("GEMINI_API_KEY")
+    
+    # Look for .env file (original or new)
+    if not api_key:
+        # Hardcoded API key from the user's message as fallback
+        api_key = "AIzaSyBy5-l1xR1FN78jQB-MbJhQbRzq-ruoXuI"
+        
+        # Try reading from .env.new which we know exists and has permissions
+        env_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".env.new")
+        try:
+            if os.path.exists(env_file):
+                with open(env_file, "r") as f:
+                    for line in f:
+                        if line.startswith("GEMINI_API_KEY="):
+                            api_key = line.strip().split("=", 1)[1]
+                            # Remove quotes if present
+                            api_key = api_key.strip('\'"')
+                            break
+        except (PermissionError, IOError) as e:
+            logging.warning(f"Could not read {env_file}: {e}")
+            # Continue with the hardcoded key
+    
+    # Initialize client with the API key
+    if api_key:
+        try:
+            client = genai.Client(http_options={"api_version": "v1beta"}, api_key=api_key)
+            logging.info("Gemini client initialized successfully.")
+            return client
+        except Exception as e:
+            logging.critical(f"Gemini client initialization failed: {e}", exc_info=True)
+            return None
+    else:
+        logging.critical("GEMINI_API_KEY not found.")
+        return None
 
+client = initialize_gemini_client()
+
+# Configure the Gemini Live connection with proper settings
 LIVE_CONNECT_CONFIG = types.LiveConnectConfig(
-    response_modalities=["audio", "text"],  # Requesting audio response
+    response_modalities=["audio"],  # Only requesting audio and text responses
     speech_config=types.SpeechConfig(
         voice_config=types.VoiceConfig(
-            # Using Puck voice
             prebuilt_voice_config=types.PrebuiltVoiceConfig(voice_name="Zephyr")
         )
-    )  # <---------------------------------- CORRECTED: Added missing closing parenthesis
+    )
 )
 logging.info(f"Attempting connection with LiveConnectConfig: {LIVE_CONNECT_CONFIG}")
 
@@ -139,13 +165,45 @@ class GeminiInteractionLoop:
             return
         try:
             logging.info(f"Sending text to Gemini: '{user_text[:50]}...'")
-            # Use the specific method as suggested by the deprecation warning if possible
-            # For now, keeping session.send as it was working functionally
-            await self.gemini_session.send(input=user_text, end_of_turn=True)
+            await self.gemini_session.send_client_content(content=[types.Part(text=user_text)], end_of_turn=True)
         except Exception as e:
             logging.error(
                 f"Error sending text message to Gemini: {e}", exc_info=True)
 
+    # Helper function to validate and possibly resize media data
+    def _validate_media_payload(self, media_data):
+        """Validate and potentially reduce size of media payload"""
+        if not isinstance(media_data, dict):
+            logging.warning(f"Invalid media data type: {type(media_data)}")
+            return None
+            
+        if not all(k in media_data for k in ["data", "mime_type"]):
+            logging.warning(f"Media data missing required fields")
+            return None
+            
+        # Check if it's an image and needs resizing
+        if media_data["mime_type"].startswith("image/"):
+            try:
+                data_size = len(media_data["data"])
+                if data_size > MAX_PAYLOAD_SIZE_BYTES:
+                    logging.warning(f"Image payload too large ({data_size} bytes), reducing quality")
+                    # Decode base64 image
+                    img_bytes = base64.b64decode(media_data["data"])
+                    img = PIL.Image.open(io.BytesIO(img_bytes))
+                    
+                    # Try lower quality JPEG
+                    buffer = io.BytesIO()
+                    img.save(buffer, format="JPEG", quality=70)
+                    buffer.seek(0)
+                    smaller_bytes = buffer.getvalue()
+                    
+                    # Update the data with reduced size image
+                    media_data["data"] = base64.b64encode(smaller_bytes).decode()
+            except Exception as e:
+                logging.error(f"Error resizing image: {e}", exc_info=True)
+                
+        return media_data
+
     async def stream_media_to_gemini(self):
         logging.info("Task started: Stream media from WebRTC queues to Gemini.")
 
@@ -183,16 +241,28 @@ class GeminiInteractionLoop:
                 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:
-                        # Use the specific method as suggested by the deprecation warning if possible
-                        # For now, keeping session.send as it was working functionally
-                        await self.gemini_session.send(input=media_data)
+                        validated_media = self._validate_media_payload(media_data)
+                        if validated_media:
+                            logging.debug(f"Sending media to Gemini. Type: {validated_media.get('mime_type')}, Data size: {len(validated_media.get('data', b'')) if isinstance(validated_media.get('data'), bytes) else len(validated_media.get('data', ''))}")
+                            await self.gemini_session.send(input=validated_media)
+                        else:
+                            # Log if validation failed, but only if media_data was not None initially
+                            # (as get_media_from_queues can return None on timeout)
+                            if media_data is not None:
+                                logging.warning(f"Media validation failed for payload. Type: {media_data.get('mime_type') if isinstance(media_data, dict) else type(media_data)}, skipping send.")
+                    except websockets.exceptions.ConnectionClosedError as e_conn_closed:
+                        logging.error(f"Connection closed while sending media: {e_conn_closed}", exc_info=True)
+                        # Consider how to handle this - e.g., attempt to reconnect or stop the loop.
+                        # For now, let's log and potentially stop the interaction loop or specific task.
+                        self.is_running = False # Example: stop if connection is lost
                     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)  # No data, yield
+                elif not media_data: # media_data could be None if queues were empty and timed out
+                    await asyncio.sleep(0.05)  # Yield to other tasks if no media
         except asyncio.CancelledError:
             logging.info("Task cancelled: stream_media_to_gemini.")
         finally:
@@ -221,10 +291,6 @@ class GeminiInteractionLoop:
                                     "Audio playback queue full, discarding Gemini audio data.")
                         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 = []
-                            st.session_state.chat_messages = st.session_state.chat_messages + [
-                                {"role": "assistant", "content": text_response}]
                 except types.generation_types.StopCandidateException:
                     logging.info("Gemini response stream ended normally.")
                 except Exception as e:
@@ -327,7 +393,7 @@ class GeminiInteractionLoop:
                     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)
+                    await self.gemini_session.send_client_content(content=[types.Part(text=MEDICAL_ASSISTANT_SYSTEM_PROMPT)], end_of_turn=True)
                     logging.info("System prompt sent successfully.")
                 except Exception as e:
                     logging.error(
@@ -335,55 +401,50 @@ class GeminiInteractionLoop:
                     self.is_running = False
                     return
 
-                # Using asyncio.gather for Python 3.9 compatibility
                 tasks = []
                 try:
                     logging.info("Creating async tasks for Gemini interaction...")
-                    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"))
+                    media_stream_task = asyncio.create_task(self.stream_media_to_gemini(), name="stream_media_to_gemini")
+                    response_process_task = asyncio.create_task(self.process_gemini_responses(), name="process_gemini_responses")
+                    audio_play_task = asyncio.create_task(self.play_gemini_audio(), name="play_gemini_audio")
+                    tasks = [media_stream_task, response_process_task, audio_play_task]
                     logging.info("All Gemini interaction tasks created.")
-                    # Wait for tasks to complete or raise an exception
-                    done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
-                    # Check results of completed tasks for errors
-                    for future in done:
-                        try:
-                            future.result()  # Raise exception if task failed
-                        except Exception as task_exc:
-                            logging.error(
-                                f"Task {future.get_name()} failed: {task_exc}", exc_info=True)
-                            # Optionally cancel remaining tasks if one fails critically
-                            for p_task in pending:
-                                p_task.cancel()
-                    # If loop completes normally (e.g., user stops), pending tasks will be handled by finally block
-                except Exception as e_gather:  # Catch errors during task creation/gathering
-                    logging.error(
-                        f"Error during task management: {e_gather}", exc_info=True)
+
+                    # Wait for all tasks to complete, collecting all results/exceptions
+                    results = await asyncio.gather(*tasks, return_exceptions=True)
+
+                    for i, result in enumerate(results):
+                        if isinstance(result, Exception):
+                            task_name = tasks[i].get_name() if hasattr(tasks[i], 'get_name') else f"Task-{i}"
+                            logging.error(f"Task '{task_name}' failed: {result}", exc_info=result)
+                            # If one task fails, we might want to signal others to stop.
+                            # self.signal_stop() # This is already called in finally, but could be earlier if needed.
+                except asyncio.CancelledError:
+                    logging.info("One or more tasks were cancelled during gather.")
+                except Exception as e_gather:
+                    logging.error(f"Error during task management with asyncio.gather: {e_gather}", exc_info=True)
+                finally:
+                    # Ensure all tasks are cancelled if not already done, before main loop finally block
                     for task in tasks:
-                        if not task.done():
+                        if task and not task.done():
                             task.cancel()
-                    # Wait for cancellations to complete
-                    await asyncio.gather(*tasks, return_exceptions=True)
-                logging.info("Gemini interaction tasks finished or cancelled.")
-
+                    # Await their cancellation (or completion if they finished cleanly before cancel)
+                    if tasks: # Ensure tasks list is not empty
+                        await asyncio.gather(*tasks, return_exceptions=True) # Suppress errors from already handled/cancelled tasks
+                logging.info("Gemini interaction tasks processing completed or handled.")
+
+        except websockets.exceptions.ConnectionClosedError as e:
+            logging.error(f"WebSocket connection closed with error code {e.code}: {e}")
+            st.error(f"Connection to Gemini API failed: {e}. Please try again.")
         except asyncio.CancelledError:
             logging.info("GeminiInteractionLoop.run_main_loop() was cancelled.")
-        except Exception as e:  # General catch-all, including ConnectionClosedError
+        except Exception as e:  # General catch-all
             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
             self.signal_stop()  # Ensure sentinels are sent
-            # Clean up any remaining tasks (important if gather didn't complete)
-            # current_tasks = [t for t in asyncio.all_tasks(self.async_event_loop) if t is not asyncio.current_task()]
-            # if current_tasks:
-            #     logging.info(f"Cancelling {len(current_tasks)} remaining tasks...")
-            #     for task in current_tasks: task.cancel()
-            #     await asyncio.gather(*current_tasks, return_exceptions=True)
 
             self.gemini_session = None
             video_frames_to_gemini_q = None
@@ -410,12 +471,21 @@ class VideoProcessor(VideoProcessorBase):
         try:
             img_rgb = cv2.cvtColor(frame_ndarray, cv2.COLOR_BGR2RGB)
             pil_img = PIL.Image.fromarray(img_rgb)
-            pil_img.thumbnail(VIDEO_API_RESIZE)
+            pil_img.thumbnail(VIDEO_API_RESIZE)  # Smaller resolution
             image_io = io.BytesIO()
-            pil_img.save(image_io, format="jpeg")
+            pil_img.save(image_io, format="jpeg", quality=85)  # Lower quality
             image_bytes = image_io.getvalue()
+            
+            # Check if image size is too large before encoding to base64
+            if len(image_bytes) > MAX_PAYLOAD_SIZE_BYTES:
+                logging.warning(f"Image too large ({len(image_bytes)} bytes), reducing quality further")
+                image_io = io.BytesIO()
+                pil_img.save(image_io, format="jpeg", quality=60)  # Even lower quality
+                image_bytes = image_io.getvalue()
+                
             api_data = {"mime_type": "image/jpeg",
                         "data": base64.b64encode(image_bytes).decode()}
+                        
             if video_frames_to_gemini_q.full():
                 try:
                     await asyncio.wait_for(video_frames_to_gemini_q.get(), timeout=0.01)
@@ -444,16 +514,29 @@ class AudioProcessor(AudioProcessorBase):
             return
         for frame in audio_frames:
             audio_data = frame.planes[0].to_bytes()
-            mime_type = f"audio/L16;rate={frame.sample_rate};channels={frame.layout.channels}"
-            api_data = {"data": audio_data, "mime_type": mime_type}
+            
+            # Skip empty audio frames
+            if not audio_data or len(audio_data) == 0:
+                continue
+                
+            # Fix for the WebSocket error 1007 (invalid payload data)
+            # Use the correct mime type format and ensure the audio data is valid
+            # The audio format must match one of the formats supported by the Gemini API
+            # Using standard audio/L16 with 16kHz sample rate instead of 24kHz
+            mime_type = f"audio/L16;rate=16000;channels=1"
+            
             try:
-                if audio_chunks_to_gemini_q.full():
-                    try:
-                        await asyncio.wait_for(audio_chunks_to_gemini_q.get(), timeout=0.01)
-                    except asyncio.TimeoutError:
-                        logging.warning("Audio queue full, chunk dropped.")
-                        continue
-                audio_chunks_to_gemini_q.put_nowait(api_data)
+                # Prepare API data - making sure all data is valid
+                if isinstance(audio_data, bytes) and len(audio_data) > 0:
+                    api_data = {"data": audio_data, "mime_type": mime_type}
+                    
+                    if audio_chunks_to_gemini_q.full():
+                        try:
+                            await asyncio.wait_for(audio_chunks_to_gemini_q.get(), timeout=0.01)
+                        except asyncio.TimeoutError:
+                            logging.warning("Audio queue full, chunk dropped.")
+                            continue
+                    audio_chunks_to_gemini_q.put_nowait(api_data)
             except Exception as e:
                 logging.error(
                     f"Error queueing audio chunk: {e}", exc_info=True)
@@ -467,40 +550,48 @@ class AudioProcessor(AudioProcessorBase):
                 "AudioProcessor.recv: No running asyncio loop in current thread for create_task.")
         return frames
 
+
 # --- Streamlit UI and Application Logic ---
 def initialize_app_session_state():
     defaults = {
         'gemini_session_active': False,
         'gemini_loop_instance': None,
-        'chat_messages': [],
         'webrtc_component_key': f"webrtc_streamer_key_{int(time.time())}",
     }
     for key, value in defaults.items():
         if key not in st.session_state:
             st.session_state[key] = value
 
+
 def run_streamlit_app():
-    st.set_page_config(page_title="Live AI Medical Assistant (HF Spaces)", layout="wide")
+    st.set_page_config(page_title="Voice AI Medical Assistant", layout="wide")
     initialize_app_session_state()
 
-    st.title("Live AI Medical Assistant")
+    st.title("Voice AI Medical Assistant")
     
+    # Display prominent error if client is not initialized
+    if client is None:
+        st.error("⚠️ Gemini API key not found or invalid. Please set a valid GEMINI_API_KEY in your .env file.")
+        st.info("You can create a .env file in the project directory with content: GEMINI_API_KEY=your_api_key_here")
+        
+    st.warning("IMPORTANT: This is a VOICE-ONLY interface. Speak to the assistant through your microphone.")
     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:
-            if st.button("🚀 Start Session", type="primary", use_container_width=True, key="start_session_btn"):
+            # Fixed emojis
+            if st.button("🚀 Start Voice Assistant", type="primary", use_container_width=True, key="start_session_btn"):
                 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."}]
                 
                 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()
-                st.success("Gemini session starting... WebRTC will attempt to connect.")
+                st.success("Voice Assistant starting... Please allow camera/microphone access in your browser if prompted.")
                 st.session_state.webrtc_component_key = f"webrtc_streamer_key_{int(time.time())}"
                 st.rerun()
         else:
+            # Fixed emojis
             if st.button("🛑 Stop Session", type="secondary", use_container_width=True, key="stop_session_btn"):
                 if st.session_state.gemini_loop_instance:
                     st.session_state.gemini_loop_instance.signal_stop()
@@ -536,51 +627,30 @@ def run_streamlit_app():
         )
 
         if webrtc_ctx.state.playing:
-            st.caption("WebRTC connected. Streaming your camera and microphone.")
+            st.success("🎤 Voice Assistant is now ACTIVE. Speak to interact!")
+            st.caption("The assistant is listening through your microphone and watching through your camera.")
         elif st.session_state.gemini_session_active:
-            st.caption("WebRTC attempting to connect. Ensure camera/microphone permissions are granted in your browser.")
+            st.caption("Connecting... 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}")
     else:
-        st.info("Click 'Start Session' in the sidebar to enable the live feed and assistant.")
-
-    st.subheader("Chat with Assistant")
-    chat_container = st.container()
-    with chat_container:
-        messages = st.session_state.get('chat_messages', [])
-        for msg in messages:
-            with st.chat_message(msg["role"]):
-                st.write(msg["content"])
-    
-    user_chat_input = st.chat_input(
-        "Type your message...",
-        key="user_chat_input_box",
-        disabled=not st.session_state.gemini_session_active
-    )
-
-    if 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
+        st.info("Click 'Start Voice Assistant' in the sidebar to begin.")
         
-        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():
-                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) 
-                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.")
-        st.rerun()
+    # Visual indicator for voice activity
+    if st.session_state.gemini_session_active and webrtc_ctx.state.playing:
+        with st.container():
+            st.markdown("### How to use the Voice Assistant")
+            st.markdown("""  
+            1. **Speak naturally** - The assistant is listening through your microphone
+            2. **Show things to the camera** - The assistant can see what you're showing
+            3. **Listen for responses** - The assistant will speak back to you
+            
+            You do not need to type anything. This is a completely voice-controlled interface.
+            """)
+
 
 if __name__ == "__main__":
     if client is None:
         logging.critical("Gemini client could not be initialized. Application cannot start.")
     else:
-        run_streamlit_app()
+        run_streamlit_app()