Update app.py

app.py

@@ -59,15 +59,25 @@ class GeminiHandler(StreamHandler):
         )
 
     def _initialize_websocket(self):
-        self.ws = websockets.sync.client.connect(self.config.ws_url)
-        initial_request = {
-            'setup': {
-                'model': self.config.model,
+        try:
+            self.ws = websockets.sync.client.connect(
+                self.config.ws_url,
+                timeout=30
+            )
+            initial_request = {
+                'setup': {
+                    'model': self.config.model,
+                }
             }
-        }
-        self.ws.send(json.dumps(initial_request))
-        setup_response = json.loads(self.ws.recv())
-        print(f"Setup response: {setup_response}")
+            self.ws.send(json.dumps(initial_request))
+            setup_response = json.loads(self.ws.recv())
+            print(f"Setup response: {setup_response}")
+        except websockets.exceptions.WebSocketException as e:
+            print(f"WebSocket connection failed: {str(e)}")
+            self.ws = None
+        except Exception as e:
+            print(f"Setup failed: {str(e)}")
+            self.ws = None
 
     def receive(self, frame: tuple[int, np.ndarray]) -> None:
         try:
@@ -102,15 +112,23 @@ class GeminiHandler(StreamHandler):
     def generator(self):
         while True:
             if not self.ws:
+                print("WebSocket not connected")
                 yield None
                 continue
 
-            message = self.ws.recv()
-            msg = json.loads(message)
-            
-            if 'serverContent' in msg:
-                content = msg['serverContent'].get('modelTurn', {})
-                yield from self._process_server_content(content)
+            try:
+                message = self.ws.recv(timeout=5)
+                msg = json.loads(message)
+                
+                if 'serverContent' in msg:
+                    content = msg['serverContent'].get('modelTurn', {})
+                    yield from self._process_server_content(content)
+            except websockets.exceptions.TimeoutError:
+                print("Timeout waiting for server response")
+                yield None
+            except Exception as e:
+                print(f"Error in generator: {str(e)}")
+                yield None
 
     def emit(self) -> tuple[int, np.ndarray] | None:
         if not self.ws:
@@ -132,6 +150,15 @@ class GeminiHandler(StreamHandler):
         if self.ws:
             self.ws.close()
 
+    def check_connection(self):
+        try:
+            if not self.ws or self.ws.closed:
+                self._initialize_websocket()
+            return True
+        except Exception as e:
+            print(f"Connection check failed: {str(e)}")
+            return False
+
 class GeminiVoiceChat:
     def __init__(self):
         self.demo = self._create_interface()
@@ -145,6 +172,8 @@ class GeminiVoiceChat:
                 </div>
             """)
             
+            error_message = gr.Textbox(label="Status", interactive=False)
+            
             webrtc = WebRTC(
                 label="Conversation",
                 modality="audio",
@@ -154,14 +183,18 @@ class GeminiVoiceChat:
             webrtc.stream(
                 GeminiHandler(),
                 inputs=[webrtc],
-                outputs=[webrtc],
+                outputs=[webrtc, error_message],
                 time_limit=90,
                 concurrency_limit=10
             )
         return demo
 
     def launch(self):
-        self.demo.launch()
+        self.demo.launch(
+            server_name="0.0.0.0",
+            server_port=int(os.environ.get("PORT", 7860)),
+            share=True  # Only if you want to create a public URL
+        )
 
 if __name__ == "__main__":
     app = GeminiVoiceChat()