Create app.py

app.py

@@ -0,0 +1,144 @@
+import gradio as gr
+from gradio_webrtc import WebRTC, StreamHandler
+import websockets.sync.client
+import numpy as np
+import json
+import base64
+import os
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# Get API key and validate
+API_KEY = os.getenv('GOOGLE_API_KEY')
+if not API_KEY:
+    raise ValueError("GOOGLE_API_KEY not found in environment variables. Please set it in your .env file.")
+
+class GeminiHandler(StreamHandler):
+    def __init__(self,
+                 expected_layout="mono",
+                 output_sample_rate=24000,
+                 output_frame_size=480) -> None:
+        self.ws = None
+        self.HOST = 'generativelanguage.googleapis.com'
+        self.API_KEY = API_KEY  # Use the validated API key
+        self.MODEL = 'models/gemini-2.0-flash-exp'
+        self.ws_url = f'wss://{self.HOST}/ws/google.ai.generativelanguage.v1alpha.GenerativeService.BidiGenerateContent?key={self.API_KEY}'
+        self.all_output_data = None
+        super().__init__(expected_layout, output_sample_rate, output_frame_size,
+                        input_sample_rate=24000)
+
+    def copy(self):
+        return GeminiHandler(
+            expected_layout=self.expected_layout,
+            output_sample_rate=self.output_sample_rate,
+            output_frame_size=self.output_frame_size
+        )
+
+    def encode_audio_input(self, data):
+        encoded = base64.b64encode(data.tobytes()).decode('UTF-8')
+        return {
+            'realtimeInput': {
+                'mediaChunks': [{
+                    'mimeType': f'audio/pcm;rate={self.output_sample_rate}',
+                    'data': encoded,
+                }],
+            },
+        }
+
+    def receive(self, frame: tuple[int, np.ndarray]) -> None:
+        try:
+            if not self.ws:
+                self.ws = websockets.sync.client.connect(self.ws_url)
+                initial_request = {
+                    'setup': {
+                        'model': self.MODEL,
+                    }
+                }
+                self.ws.send(json.dumps(initial_request))
+                setup_response = json.loads(self.ws.recv())
+                print(f"Setup response: {setup_response}")
+
+            _, array = frame
+            array = array.squeeze()
+            # Send audio data
+            audio_message = self.encode_audio_input(array)
+            self.ws.send(json.dumps(audio_message))
+        except Exception as e:
+            print(f"Error in receive: {str(e)}")
+            if self.ws:
+                self.ws.close()
+            self.ws = None
+
+    def generator(self):
+        while True:
+            if not self.ws:
+                yield None
+                continue
+
+            message = self.ws.recv()
+            msg = json.loads(message)
+            
+            if 'serverContent' in msg:
+                content = msg['serverContent'].get('modelTurn', {})
+                for part in content.get('parts', []):
+                    data = part.get('inlineData', {}).get('data', '')
+                    if data:
+                        audio_data = base64.b64decode(data)
+                        audio_array = np.frombuffer(audio_data, dtype=np.int16)
+                        if self.all_output_data is None:
+                            self.all_output_data = audio_array
+                        else:
+                            self.all_output_data = np.concatenate((self.all_output_data, audio_array))
+                        
+                        while self.all_output_data.shape[-1] >= self.output_frame_size:
+                            yield (self.output_sample_rate, 
+                                  self.all_output_data[:self.output_frame_size].reshape(1, -1))
+                            self.all_output_data = self.all_output_data[self.output_frame_size:]
+
+    def emit(self) -> tuple[int, np.ndarray] | None:
+        if not self.ws:
+            return None
+        if not hasattr(self, '_generator'):
+            self._generator = self.generator()
+        try:
+            return next(self._generator)
+        except StopIteration:
+            self.reset()
+            return None
+
+    def reset(self) -> None:
+        if hasattr(self, '_generator'):
+            delattr(self, '_generator')
+        self.all_output_data = None
+
+    def shutdown(self) -> None:
+        if self.ws:
+            self.ws.close()
+
+# Create Gradio interface
+with gr.Blocks() as demo:
+    gr.HTML("""
+        <div style='text-align: center'>
+            <h1>Gemini 2.0 Voice Chat</h1>
+            <p>Speak with Gemini using real-time audio streaming</p>
+        </div>
+    """)
+    
+    chatbot = gr.Chatbot(type='messages')
+    webrtc = WebRTC(
+        label="Conversation",
+        modality="audio",
+        mode="send-receive"
+    )
+    
+    webrtc.stream(
+        GeminiHandler(),
+        inputs=[webrtc],
+        outputs=[webrtc],
+        time_limit=90
+    )
+
+if __name__ == "__main__":
+    demo.launch()