# -*- coding: utf-8 -*- import streamlit as st import os import asyncio import base64 import io import threading import traceback import atexit # Correctly imported import time import logging from dotenv import load_dotenv import cv2 # For image processing import pyaudio # For audio PLAYBACK import PIL.Image from google import genai from google.genai import types # streamlit-webrtc components from streamlit_webrtc import ( webrtc_streamer, WebRtcMode, AudioProcessorBase, VideoProcessorBase, ) # from aiortc import RTCIceServer, RTCConfiguration # Not needed directly # --- Configuration --- load_dotenv() logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(threadName)s - %(levelname)s - %(message)s') # Audio configuration PYAUDIO_FORMAT = pyaudio.paInt16 PYAUDIO_CHANNELS = 1 WEBRTC_REQUESTED_AUDIO_CHANNELS = 1 WEBRTC_REQUESTED_SEND_SAMPLE_RATE = 16000 GEMINI_AUDIO_RECEIVE_SAMPLE_RATE = 24000 PYAUDIO_PLAYBACK_CHUNK_SIZE = 1024 AUDIO_PLAYBACK_QUEUE_MAXSIZE = 50 MEDIA_TO_GEMINI_QUEUE_MAXSIZE = 30 # Video configuration VIDEO_FPS_TO_GEMINI = 2 VIDEO_API_RESIZE = (1024, 1024) # !!! IMPORTANT: Verify this model name is correct for the Live API !!! MODEL_NAME = "models/gemini-2.0-flash-live-001" 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. Your responsibilities are: 1. **Visual Observation and Description:** Carefully examine the images or video feed. Describe relevant details you observe. 2. **General Information (Non-Diagnostic):** Provide general information related to what is visually presented, if applicable. You are not a diagnostic tool. 3. **Safety and Disclaimer (CRITICAL):** * You are an AI assistant, **NOT a medical doctor or a substitute for one.** * **DO NOT provide medical diagnoses, treatment advice, or interpret medical results (e.g., X-rays, scans, lab reports).** * When appropriate, and always if the user seems to be seeking diagnosis or treatment, explicitly state your limitations and **strongly advise the user to consult a qualified healthcare professional.** * If you see something that *appears* visually concerning (e.g., an unusual skin lesion, signs of injury), you may gently suggest it might be wise to have it looked at by a professional, without speculating on what it is. 4. **Tone:** Maintain a helpful, empathetic, and calm tone. 5. **Interaction:** After this initial instruction, you can make a brief acknowledgment of your role (e.g., "I'm ready to assist by looking at what you show me. Please remember to consult a doctor for medical advice."). Then, focus on responding to the user's visual input and questions. Example of a disclaimer you might use: "As an AI assistant, I can describe what I see, but I can't provide medical advice or diagnoses. For any health concerns, it's always best to speak with a doctor or other healthcare professional." """ # --- PyAudio Global Instance and Cleanup --- pya = pyaudio.PyAudio() def cleanup_pyaudio(): logging.info("Terminating PyAudio instance.") if pya: pya.terminate() atexit.register(cleanup_pyaudio) # --- Global Queues - Declare as None, initialize later --- video_frames_to_gemini_q: asyncio.Queue = None 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) 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() # 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. 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")) # ), ) logging.info(f"Using LiveConnectConfig: {LIVE_CONNECT_CONFIG}") # Log the config being used # --- Backend Gemini Interaction Loop --- class GeminiInteractionLoop: def __init__(self): self.gemini_session = None self.async_event_loop = None self.is_running = True self.playback_stream = None async def send_text_input_to_gemini(self, user_text): if not user_text or not self.gemini_session or not self.is_running: logging.warning("Cannot send text. Session not active, no text, or not running.") return try: logging.info(f"Sending text to Gemini: '{user_text[:50]}...'") await self.gemini_session.send(input=user_text, end_of_turn=True) except Exception as e: logging.error(f"Error sending text message to Gemini: {e}", exc_info=True) async def stream_media_to_gemini(self): logging.info("Task started: Stream media from WebRTC queues to Gemini.") async def get_media_from_queues(): if video_frames_to_gemini_q is None or audio_chunks_to_gemini_q is None: await asyncio.sleep(0.1); return None try: video_frame = await asyncio.wait_for(video_frames_to_gemini_q.get(), timeout=0.02) 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) 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 try: 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 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) except asyncio.CancelledError: logging.info("Task cancelled: stream_media_to_gemini.") finally: logging.info("Task finished: stream_media_to_gemini.") async def process_gemini_responses(self): logging.info("Task started: Process responses from Gemini.") try: while self.is_running: if not self.gemini_session: await asyncio.sleep(0.1); continue if audio_from_gemini_playback_q is None: await asyncio.sleep(0.1); continue try: turn_response = self.gemini_session.receive() async for chunk in turn_response: if not self.is_running: break if audio_data := chunk.data: if not audio_from_gemini_playback_q.full(): audio_from_gemini_playback_q.put_nowait(audio_data) else: logging.warning("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 = [] # 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}] 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) await asyncio.sleep(0.1) except asyncio.CancelledError: logging.info("Task cancelled: process_gemini_responses.") finally: logging.info("Task finished: process_gemini_responses.") 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. try: while audio_from_gemini_playback_q is None and self.is_running: await asyncio.sleep(0.1) if not self.is_running: return self.playback_stream = await asyncio.to_thread( pya.open, format=PYAUDIO_FORMAT, channels=PYAUDIO_CHANNELS, rate=GEMINI_AUDIO_RECEIVE_SAMPLE_RATE, output=True, frames_per_buffer=PYAUDIO_PLAYBACK_CHUNK_SIZE ) logging.info(f"PyAudio playback stream opened at {GEMINI_AUDIO_RECEIVE_SAMPLE_RATE} Hz.") while self.is_running: try: audio_chunk = await asyncio.wait_for(audio_from_gemini_playback_q.get(), timeout=1.0) 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 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) 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) async def run_main_loop(self): global video_frames_to_gemini_q, audio_chunks_to_gemini_q, audio_from_gemini_playback_q self.async_event_loop = asyncio.get_running_loop() self.is_running = True logging.info("GeminiInteractionLoop run_main_loop starting...") video_frames_to_gemini_q = asyncio.Queue(maxsize=MEDIA_TO_GEMINI_QUEUE_MAXSIZE) audio_chunks_to_gemini_q = asyncio.Queue(maxsize=MEDIA_TO_GEMINI_QUEUE_MAXSIZE) audio_from_gemini_playback_q = asyncio.Queue(maxsize=AUDIO_PLAYBACK_QUEUE_MAXSIZE) logging.info("Asyncio queues initialized in 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.") try: logging.info("Sending system prompt to Gemini...") await self.gemini_session.send(input=MEDICAL_ASSISTANT_SYSTEM_PROMPT, end_of_turn=False) logging.info("System prompt sent successfully.") except Exception as e: 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 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.") 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 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 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.") # --- WebRTC Media Processors --- class VideoProcessor(VideoProcessorBase): def __init__(self): self.frame_counter = 0 self.last_gemini_send_time = time.monotonic() async def _process_and_queue_frame_async(self, frame_ndarray): if video_frames_to_gemini_q is None: return self.frame_counter += 1 current_time = time.monotonic() if (current_time - self.last_gemini_send_time) < (1.0 / VIDEO_FPS_TO_GEMINI): return self.last_gemini_send_time = current_time try: img_rgb = cv2.cvtColor(frame_ndarray, cv2.COLOR_BGR2RGB) pil_img = PIL.Image.fromarray(img_rgb) pil_img.thumbnail(VIDEO_API_RESIZE) image_io = io.BytesIO() pil_img.save(image_io, format="jpeg") 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) except asyncio.TimeoutError: logging.warning("Video queue full, frame dropped."); return video_frames_to_gemini_q.put_nowait(api_data) except Exception as e: logging.error(f"Error processing/queueing video frame: {e}", exc_info=True) async def recv(self, frame): img_bgr = frame.to_ndarray(format="bgr24") try: loop = asyncio.get_running_loop() loop.create_task(self._process_and_queue_frame_async(img_bgr)) except RuntimeError: logging.error("VideoProcessor.recv: No running asyncio loop in current thread for create_task.") return frame class AudioProcessor(AudioProcessorBase): async def _process_and_queue_audio_async(self, audio_frames): if audio_chunks_to_gemini_q is None: 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} 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) except Exception as e: logging.error(f"Error queueing audio chunk: {e}", exc_info=True) async def recv(self, frames): try: loop = asyncio.get_running_loop() loop.create_task(self._process_and_queue_audio_async(frames)) except RuntimeError: logging.error("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") initialize_app_session_state() st.title("Live AI Medical Assistant") 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: if st.button("🚀 Start Session", 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."}] # 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() st.success("Gemini session starting... WebRTC will attempt to connect.") st.session_state.webrtc_component_key = f"webrtc_streamer_key_{int(time.time())}" st.rerun() else: 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() st.session_state.gemini_loop_instance = None st.session_state.gemini_session_active = False st.warning("Session stopping...") time.sleep(0.5) st.rerun() if st.session_state.gemini_session_active: st.subheader("Your Live Feed (from your browser)") MEDIA_STREAM_CONSTRAINTS = { "video": True, "audio": { "sampleRate": {"ideal": WEBRTC_REQUESTED_SEND_SAMPLE_RATE}, "channelCount": {"exact": WEBRTC_REQUESTED_AUDIO_CHANNELS}, "echoCancellation": True, "noiseSuppression": True } } webrtc_ctx = webrtc_streamer( key=st.session_state.webrtc_component_key, mode=WebRtcMode.SENDONLY, rtc_configuration={ "iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}] }, media_stream_constraints=MEDIA_STREAM_CONSTRAINTS, video_processor_factory=VideoProcessor, audio_processor_factory=AudioProcessor, async_processing=True, ) if webrtc_ctx.state.playing: st.caption("WebRTC connected. Streaming your camera and microphone.") elif st.session_state.gemini_session_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}") else: 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", disabled=not st.session_state.gemini_session_active ) 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}] # 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) 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()