Update app.py

app.py

@@ -1,193 +1,476 @@
-import asyncio
-import base64
+"""
+Real-Time Screen Assistant - Refactored for Google GenAI Live API + FastRTC
+
+This application transforms the original screenshot analyzer into a real-time
+screen sharing assistant with voice interaction, following the refactoring 
+instructions for live streaming capabilities.
+"""
+
 import os
+import asyncio
 import time
-from io import BytesIO
-
-import gradio as gr
 import numpy as np
-import websockets
-from dotenv import load_dotenv
-from fastrtc import (
-    AsyncAudioVideoStreamHandler,
-    Stream,
-    WebRTC,
-    get_cloudflare_turn_credentials_async,
-    wait_for_item,
-)
+import numpy.typing as npt
+import cv2
+import gradio as gr
+from fastrtc import Stream, AsyncAudioVideoStreamHandler, get_cloudflare_turn_credentials_async, ReplyOnPause
 from google import genai
-from gradio.utils import get_space
-from PIL import Image
-
-load_dotenv()
-
-
-def encode_audio(data: np.ndarray) -> dict:
-    """Encode Audio data to send to the server"""
-    return {
-        "mime_type": "audio/pcm",
-        "data": base64.b64encode(data.tobytes()).decode("UTF-8"),
-    }
+from google.genai import types
 
+# Environment variable for API key
+API_KEY = os.getenv("GEMINI_API_KEY", "")
 
-def encode_image(data: np.ndarray) -> dict:
-    with BytesIO() as output_bytes:
-        pil_image = Image.fromarray(data)
-        pil_image.save(output_bytes, "JPEG")
-        bytes_data = output_bytes.getvalue()
-    base64_str = str(base64.b64encode(bytes_data), "utf-8")
-    return {"mime_type": "image/jpeg", "data": base64_str}
-
-
-class GeminiHandler(AsyncAudioVideoStreamHandler):
-    def __init__(
-        self,
-    ) -> None:
+class RealTimeScreenAssistant(AsyncAudioVideoStreamHandler):
+    """
+    Real-time screen assistant implementing the refactoring instructions.
+    
+    Features:
+    - Google GenAI Live API integration
+    - Real-time audio/video streaming via FastRTC
+    - Voice activity detection with ReplyOnPause
+    - Intelligent frame sampling for screen sharing
+    - Cloudflare TURN server support for HF Spaces
+    """
+    
+    def __init__(self):
         super().__init__(
-            "mono",
-            output_sample_rate=24000,
-            input_sample_rate=16000,
+            expected_layout="mono", 
+            output_sample_rate=24000, 
+            input_sample_rate=16000
         )
-        self.audio_queue = asyncio.Queue()
-        self.video_queue = asyncio.Queue()
         self.session = None
         self.last_frame_time = 0
-        self.quit = asyncio.Event()
-
-    def copy(self) -> "GeminiHandler":
-        return GeminiHandler()
-
+        self.audio_queue = asyncio.Queue()
+        self.connected = False
+        self.frame_interval = 1.0  # 1 FPS as per instructions
+        
     async def start_up(self):
-        client = genai.Client(
-            api_key=os.getenv("GEMINI_API_KEY"), http_options={"api_version": "v1alpha"}
-        )
-        config = {"response_modalities": ["AUDIO"]}
-        async with client.aio.live.connect(
-            model="gemini-2.0-flash-exp",
-            config=config,  # type: ignore
-        ) as session:
-            self.session = session
-            while not self.quit.is_set():
-                turn = self.session.receive()
+        """Initialize Google GenAI Live session as per Task 8-10"""
+        try:
+            # Re-check environment variable in case it was set after import
+            current_api_key = os.getenv("GEMINI_API_KEY", "")
+            if not current_api_key:
+                print("❌ No GEMINI_API_KEY found in environment")
+                return
+                
+            # Initialize client with v1alpha API (Task 8)
+            client = genai.Client(
+                api_key=current_api_key,
+                http_options={"api_version": "v1alpha"}
+            )
+            
+            # Configure live session (Task 9) - minimal working config
+            from google.genai.types import LiveConnectConfig
+            
+            # Start with minimal config to avoid WebSocket errors
+            config = LiveConnectConfig(
+                system_instruction=(
+                    "You are a helpful real-time assistant who watches the user's screen and provides "
+                    "guidance on using software. Give clear, step-by-step instructions based on what "
+                    "you see and hear. Be proactive in offering assistance."
+                )
+            )
+            
+            # Connect to Live API (Task 10) - using async context manager
+            self.session_context = client.aio.live.connect(
+                model="gemini-2.0-flash-live-001", 
+                config=config
+            )
+            self.session = await self.session_context.__aenter__()
+            
+            self.connected = True
+            print("✅ Connected to Google GenAI Live API")
+            
+            # Start response handler (Task 13)
+            self.response_task = asyncio.create_task(self._handle_responses())
+            
+        except Exception as e:
+            print(f"❌ Failed to connect to GenAI: {e}")
+            self.connected = False
+    
+    async def _handle_responses(self):
+        """Handle AI responses as per Task 12-13"""
+        try:
+            # The Google GenAI Live API provides an async iterator through session.receive()
+            # We need to handle this as a streaming response
+            while self.connected and self.session:
                 try:
-                    async for response in turn:
-                        if data := response.data:
-                            audio = np.frombuffer(data, dtype=np.int16).reshape(1, -1)
-                        self.audio_queue.put_nowait(audio)
-                except websockets.exceptions.ConnectionClosedOK:
-                    print("connection closed")
-                    break
-
-    async def video_receive(self, frame: np.ndarray):
-        self.video_queue.put_nowait(frame)
-
-        if self.session:
-            # send image every 1 second
-            print(time.time() - self.last_frame_time)
-            if time.time() - self.last_frame_time > 1:
-                self.last_frame_time = time.time()
-                await self.session.send(input=encode_image(frame))
-                if self.latest_args[1] is not None:
-                    await self.session.send(input=encode_image(self.latest_args[1]))
-
+                    # Get the next message from the session
+                    response_stream = self.session.receive()
+                    
+                    # Check if this is an async iterator or needs to be awaited
+                    if hasattr(response_stream, '__aiter__'):
+                        # It's an async iterator
+                        async for msg in response_stream:
+                            if not self.connected:
+                                break
+                                
+                            if msg.data:  # Audio response
+                                # Convert to numpy for FastRTC (Task 13)
+                                audio_array = np.frombuffer(msg.data, dtype=np.int16)
+                                if len(audio_array) > 0:
+                                    audio_array = audio_array.reshape(1, -1)
+                                    await self.audio_queue.put(audio_array)
+                            
+                            if msg.text:  # Text response
+                                print(f"🤖 AI: {msg.text}")
+                    else:
+                        # It's a single response that needs to be awaited
+                        msg = await response_stream
+                        if msg:
+                            if msg.data:  # Audio response
+                                audio_array = np.frombuffer(msg.data, dtype=np.int16)
+                                if len(audio_array) > 0:
+                                    audio_array = audio_array.reshape(1, -1)
+                                    await self.audio_queue.put(audio_array)
+                            
+                            if msg.text:  # Text response
+                                print(f"🤖 AI: {msg.text}")
+                                
+                except Exception as inner_e:
+                    if "connection" in str(inner_e).lower() or "closed" in str(inner_e).lower():
+                        print("🔴 Connection closed, stopping response handler")
+                        break
+                    else:
+                        print(f"⚠️  Response handling error: {inner_e}")
+                        await asyncio.sleep(0.1)  # Brief pause before retry
+                    
+        except Exception as e:
+            print(f"❌ Error handling AI responses: {e}")
+    
+    async def receive(self, frame: tuple[int, npt.NDArray[np.int16]]):
+        """Handle microphone audio (Task 11)"""
+        if not self.connected or not self.session:
+            return
+            
+        try:
+            _, audio_np = frame
+            audio_bytes = audio_np.tobytes()
+            
+            # Send audio to GenAI Live API using new non-deprecated method
+            await self.session.send_realtime_input(
+                input=types.Blob(
+                    data=audio_bytes, 
+                    mime_type="audio/pcm;rate=16000"
+                )
+            )
+        except Exception as e:
+            print(f"❌ Error sending audio: {e}")
+    
+    async def video_receive(self, frame: npt.NDArray[np.float32]):
+        """Handle screen video frames (Task 11-12)"""
+        if not self.connected or not self.session:
+            return
+            
+        try:
+            # Throttle to 1 FPS as per instructions
+            current_time = time.time()
+            if current_time - self.last_frame_time < self.frame_interval:
+                return
+                
+            self.last_frame_time = current_time
+            
+            # Convert float32 frame to uint8 for JPEG encoding
+            if frame.dtype == np.float32:
+                # Assuming frame is in range [0, 1], convert to [0, 255]
+                frame_uint8 = (frame * 255).astype(np.uint8)
+            else:
+                frame_uint8 = frame.astype(np.uint8)
+            
+            # Check for empty frame before encoding
+            if frame_uint8.size == 0 or frame_uint8.shape[0] == 0 or frame_uint8.shape[1] == 0:
+                return
+                
+            # Encode as JPEG (Task 12)
+            try:
+                success, jpg_bytes = cv2.imencode('.jpg', frame_uint8, [cv2.IMWRITE_JPEG_QUALITY, 80])
+                if not success:
+                    return
+            except cv2.error:
+                # Handle OpenCV encoding errors gracefully
+                return
+            
+            # Send to GenAI using new non-deprecated method
+            await self.session.send_realtime_input(
+                input=types.Blob(
+                    data=jpg_bytes.tobytes(), 
+                    mime_type="image/jpeg"
+                )
+            )
+            
+        except Exception as e:
+            print(f"❌ Error sending video frame: {e}")
+    
+    async def emit(self):
+        """Emit audio back to user (Task 13)"""
+        try:
+            audio_chunk = self.audio_queue.get_nowait()
+            return (24000, audio_chunk)
+        except asyncio.QueueEmpty:
+            return None
+    
+    def copy(self):
+        """Copy method required by FastRTC AsyncAudioVideoStreamHandler"""
+        # Return a new instance with same configuration
+        new_instance = RealTimeScreenAssistant()
+        new_instance.frame_interval = self.frame_interval
+        return new_instance
+    
     async def video_emit(self):
-        frame = await wait_for_item(self.video_queue, 0.01)
-        if frame is not None:
-            return frame
-        else:
-            return np.zeros((100, 100, 3), dtype=np.uint8)
-
-    async def receive(self, frame: tuple[int, np.ndarray]) -> None:
-        _, array = frame
-        array = array.squeeze()
-        audio_message = encode_audio(array)
+        """Video emit method required by FastRTC AsyncAudioVideoStreamHandler"""
+        # For this use case, we don't emit video back to user
+        return None
+    
+    async def shutdown(self):
+        """Clean shutdown (Task 17)"""
+        self.connected = False
+        
+        # Cancel response handler task if it exists
+        if hasattr(self, 'response_task') and not self.response_task.done():
+            self.response_task.cancel()
+            try:
+                await self.response_task
+            except asyncio.CancelledError:
+                pass
+        
         if self.session:
-            await self.session.send(input=audio_message)
+            try:
+                # Properly close the session using context manager
+                if hasattr(self, 'session_context'):
+                    await self.session_context.__aexit__(None, None, None)
+                else:
+                    await self.session.close()
+                print("🔴 Disconnected from GenAI Live API")
+            except Exception as e:
+                print(f"❌ Error during shutdown: {e}")
+        
+        self.session = None
+        if hasattr(self, 'session_context'):
+            self.session_context = None
 
-    async def emit(self):
-        array = await wait_for_item(self.audio_queue, 0.01)
-        if array is not None:
-            return (self.output_sample_rate, array)
-        return array
+# Global state
+app_state = {"stream": None, "handler": None, "connected": False}
 
-    async def shutdown(self) -> None:
-        if self.session:
-            self.quit.set()
-            await self.session.close()
-            self.quit.clear()
+def initialize_real_time_assistant():
+    """Initialize the real-time assistant (Task 26-29)"""
+    try:
+        # Create handler
+        handler = RealTimeScreenAssistant()
+        app_state["handler"] = handler
+        
+        # Create stream with Cloudflare TURN (Task 22-23)
+        stream = Stream(
+            handler=ReplyOnPause(handler),  # Voice activity detection (Task 3)
+            modality="audio-video", 
+            mode="send-receive",
+            rtc_configuration=get_cloudflare_turn_credentials_async,
+            time_limit=300  # 5 minute limit for Spaces
+        )
+        
+        app_state["stream"] = stream
+        return stream
+        
+    except Exception as e:
+        print(f"❌ Error creating stream: {e}")
+        return None
 
+def handle_connect():
+    """Connect button handler (Task 16)"""
+    # Re-check environment variable in case it was set after import
+    current_api_key = os.getenv("GEMINI_API_KEY", "")
+    if not current_api_key:
+        return "❌ Please set GEMINI_API_KEY environment variable"
+    
+    if app_state["connected"]:
+        return "✅ Already connected - session is active"
+    
+    app_state["connected"] = True
+    return "✅ Connecting... Please allow microphone and camera permissions"
 
-stream = Stream(
-    handler=GeminiHandler(),
-    modality="audio-video",
-    mode="send-receive",
-    rtc_configuration=get_cloudflare_turn_credentials_async if get_space() else None,
-    time_limit=180 if get_space() else None,
-    additional_inputs=[
-        gr.Image(label="Image", type="numpy", sources=["upload", "clipboard"])
-    ],
-    ui_args={
-        "icon": "https://www.gstatic.com/lamda/images/gemini_favicon_f069958c85030456e93de685481c559f160ea06b.png",
-        "pulse_color": "rgb(255, 255, 255)",
-        "icon_button_color": "rgb(255, 255, 255)",
-        "title": "Gemini Audio Video Chat",
-    },
-)
+def handle_disconnect():
+    """Disconnect button handler (Task 17)"""
+    if app_state["handler"] and app_state["connected"]:
+        asyncio.create_task(app_state["handler"].shutdown())
+        app_state["connected"] = False
+        return "🔴 Disconnected from AI assistant"
+    
+    return "Already disconnected"
 
-css = """
-#video-source {max-width: 600px !important; max-height: 600 !important;}
-"""
+# Screen sharing JavaScript (Task 16)
+screen_share_js = """
+async function shareScreen() {
+    try {
+        const screenStream = await navigator.mediaDevices.getDisplayMedia({
+            video: {
+                mediaSource: 'screen',
+                width: { ideal: 1280 },
+                height: { ideal: 720 }
+            },
+            audio: false
+        });
+        
+        const videoElements = document.querySelectorAll('video');
+        const webrtcVideo = Array.from(videoElements).find(video => 
+            video.srcObject && video.srcObject.getVideoTracks().length > 0
+        );
+        
+        if (webrtcVideo && webrtcVideo.srcObject) {
+            const videoTrack = screenStream.getVideoTracks()[0];
+            const currentTrack = webrtcVideo.srcObject.getVideoTracks()[0];
+            
+            webrtcVideo.srcObject.removeTrack(currentTrack);
+            webrtcVideo.srcObject.addTrack(videoTrack);
+            
+            videoTrack.onended = () => {
+                console.log('Screen sharing ended');
+            };
+            
+            return "🖥️ Screen sharing started";
+        } else {
+            return "❌ Could not find video element";
+        }
+        
+    } catch (error) {
+        console.error('Screen sharing error:', error);
+        if (error.name === 'NotAllowedError') {
+            return "❌ Screen sharing permission denied";
+        } else {
+            return `❌ Screen sharing failed: ${error.message}`;
+        }
+    }
+}
 
-with gr.Blocks(css=css) as demo:
-    gr.HTML(
-        """
-    <div style='display: flex; align-items: center; justify-content: center; gap: 20px'>
-        <div style="background-color: var(--block-background-fill); border-radius: 8px">
-            <img src="https://www.gstatic.com/lamda/images/gemini_favicon_f069958c85030456e93de685481c559f160ea06b.png" style="width: 100px; height: 100px;">
-        </div>
-        <div>
-            <h1>Gen AI SDK Voice Chat</h1>
-            <p>Speak with Gemini using real-time audio + video streaming</p>
-            <p>Powered by <a href="https://gradio.app/">Gradio</a> and <a href=https://freddyaboulton.github.io/gradio-webrtc/">WebRTC</a>⚡️</p>
-            <p>Get an API Key <a href="https://support.google.com/googleapi/answer/6158862?hl=en">here</a></p>
-        </div>
-    </div>
-    """
-    )
-    with gr.Row() as row:
-        with gr.Column():
-            webrtc = WebRTC(
-                label="Video Chat",
-                modality="audio-video",
-                mode="send-receive",
-                elem_id="video-source",
-                rtc_configuration=get_cloudflare_turn_credentials_async
-                if get_space()
-                else None,
-                icon="https://www.gstatic.com/lamda/images/gemini_favicon_f069958c85030456e93de685481c559f160ea06b.png",
-                pulse_color="rgb(255, 255, 255)",
-                icon_button_color="rgb(255, 255, 255)",
-            )
-        with gr.Column():
-            image_input = gr.Image(
-                label="Image", type="numpy", sources=["upload", "clipboard"]
-            )
+return shareScreen();
+"""
 
-        webrtc.stream(
-            GeminiHandler(),
-            inputs=[webrtc, image_input],
-            outputs=[webrtc],
-            time_limit=180 if get_space() else None,
-            concurrency_limit=2 if get_space() else None,
+def create_interface():
+    """Create main interface (Task 26-30)"""
+    
+    # Initialize stream
+    stream = initialize_real_time_assistant()
+    
+    with gr.Blocks(
+        title="Real-Time Screen Assistant", 
+        theme=gr.themes.Soft()
+    ) as demo:
+        
+        gr.Markdown("# 🖥️ Real-Time Screen Assistant")
+        gr.Markdown("""
+        **Live AI assistant with screen sharing and voice interaction**
+        
+        Refactored according to instructions for real-time capabilities using:
+        - Google GenAI Live API for multimodal AI processing
+        - FastRTC for low-latency audio/video streaming  
+        - Voice activity detection with ReplyOnPause
+        - Cloudflare TURN servers for HF Spaces compatibility
+        """)
+        
+        # Status display
+        status_display = gr.Textbox(
+            label="Status",
+            value="Ready to connect - Click Connect to start real-time session",
+            interactive=False
         )
+        
+        # Control buttons (Task 3, 16-17)
+        with gr.Row():
+            connect_btn = gr.Button("🔗 Connect", variant="primary")
+            screen_btn = gr.Button("🖥️ Show Your Screen", variant="secondary") 
+            disconnect_btn = gr.Button("🔴 Disconnect", variant="stop")
+        
+        # Stream interface
+        if stream and hasattr(stream, 'ui'):
+            gr.Markdown("### 📡 Live Stream")
+            # Mount the FastRTC stream UI
+            stream_ui = stream.ui
+        else:
+            stream_ui = gr.HTML("<div>⚠️ Stream initialization failed</div>")
+        
+        # Instructions (Task 1-3)
+        with gr.Accordion("📋 Instructions", open=True):
+            gr.Markdown("""
+            **How to use the real-time assistant:**
+            
+            1. **Connect**: Click Connect to start the AI session
+            2. **Permissions**: Allow microphone and camera access
+            3. **Show Screen**: Click "Show Your Screen" to share your screen
+            4. **Voice Interaction**: Simply speak - the AI will respond
+            5. **Real-time Guidance**: AI sees your screen and provides live help
+            6. **Disconnect**: Click Disconnect when finished
+            
+            **Features implemented from refactoring instructions:**
+            - ✅ FastRTC WebRTC streaming (Task 2)
+            - ✅ Google GenAI Live API integration (Task 7-15)
+            - ✅ Connect/Show Screen/Disconnect controls (Task 3, 16-17)
+            - ✅ Voice activity detection with ReplyOnPause (Task 3)
+            - ✅ Screen sharing via getDisplayMedia (Task 6)
+            - ✅ Real-time advice generation (Task 18-21)
+            - ✅ Cloudflare TURN for HF Spaces (Task 22-23)
+            """)
+        
+        # Privacy notice (Task 24-25)
+        with gr.Accordion("🔒 Privacy & Security", open=False):
+            gr.Markdown("""
+            **Privacy Notice:**
+            - Screen content and voice are processed by Google's AI services
+            - Data is transmitted securely via encrypted WebRTC connections
+            - No permanent storage - all processing is real-time
+            - You control what is shared and can disconnect anytime
+            
+            **Technical Details:**
+            - Uses Google Gemini Live API for real-time multimodal processing
+            - FastRTC provides low-latency WebRTC streaming
+            - Cloudflare TURN servers ensure reliable connectivity on HF Spaces
+            - Voice activity detection prevents interruptions
+            """)
+        
+        # Wire up controls
+        connect_btn.click(
+            fn=handle_connect,
+            outputs=[status_display]
+        )
+        
+        screen_btn.click(
+            fn=lambda: "🖥️ Requesting screen share...",
+            outputs=[status_display],
+            js=screen_share_js
+        )
+        
+        disconnect_btn.click(
+            fn=handle_disconnect,
+            outputs=[status_display]
+        )
+        
+        return demo
 
-stream.ui = demo
-
-
+# Main execution
 if __name__ == "__main__":
-    if (mode := os.getenv("MODE")) == "UI":
-        stream.ui.launch(server_port=7860)
-    elif mode == "PHONE":
-        raise ValueError("Phone mode not supported for this demo")
+    print("🖥️ Real-Time Screen Assistant")
+    print("=" * 50)
+    print("Refactored according to instructions for:")
+    print("- Google GenAI Live API integration")
+    print("- FastRTC real-time streaming") 
+    print("- Voice activity detection")
+    print("- Screen sharing capabilities")
+    print("- Cloudflare TURN for HF Spaces")
+    
+    if not API_KEY:
+        print("\n⚠️  No GEMINI_API_KEY environment variable found")
+        print("Please set your Google AI API key:")
+        print("export GEMINI_API_KEY='your-api-key-here'")
     else:
-        stream.ui.launch(server_port=7860)
+        print(f"\n✅ API key configured (length: {len(API_KEY)})")
+    
+    print("\n🚀 Starting real-time assistant...")
+    
+    try:
+        demo = create_interface()
+        demo.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=False,
+            show_error=True
+        )
+    except Exception as e:
+        print(f"❌ Failed to launch: {e}")
+        print("Ensure all dependencies are installed: pip install -r requirements.txt")