# -*- coding: utf-8 -*- import streamlit as st import os import asyncio import base64 import io import threading import traceback import atexit 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') logging.info("Application starting up...") # 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" 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. 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 = 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 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() # Gemini LiveConnectConfig - Using audio response and Puck voice as in your latest code # Ensure this configuration is valid for your API key and model. LIVE_CONNECT_CONFIG = types.LiveConnectConfig( response_modalities=["audio"], # Requesting audio response speech_config=types.SpeechConfig( voice_config=types.VoiceConfig( prebuilt_voice_config=types.PrebuiltVoiceConfig(voice_name="Puck") # Using Puck voice ) ) # <---------------------------------- CORRECTED: Added missing closing parenthesis ) logging.info(f"Attempting connection with LiveConnectConfig: {LIVE_CONNECT_CONFIG}") # --- 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]}...'") # 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) 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) 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 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 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) 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 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 = [] 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.") 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 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 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 (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) 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_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 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 try: async with client.aio.live.connect(model=MODEL_NAME, config=LIVE_CONNECT_CONFIG) as session: self.gemini_session = session 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) 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 # 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")) 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) for task in tasks: if 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.") except asyncio.CancelledError: logging.info("GeminiInteractionLoop.run_main_loop() was cancelled.") 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 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 audio_chunks_to_gemini_q = None audio_from_gemini_playback_q = None logging.info("GeminiInteractionLoop finished and global queues set to None.") # --- 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."}] 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") 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 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() if __name__ == "__main__": if client is None: logging.critical("Gemini client could not be initialized. Application cannot start.") else: run_streamlit_app()