Update streamlit_app.py

streamlit_app.py

@@ -4,6 +4,7 @@ import uuid
 import requests
 import json
 from typing import Dict, Optional
+import time
 
 class VideoProcessor(VideoProcessorBase):
     def recv(self, frame):
@@ -11,30 +12,45 @@ class VideoProcessor(VideoProcessorBase):
         return img
 
 class WebRTCConnection:
-    def __init__(self, server_url: str):
-        self.server_url = server_url
+    def __init__(self):
+        self.server_url = "http://localhost:7860"
         self.peer_id: Optional[str] = None
         
     def register_peer(self, peer_id: str) -> Dict:
-        response = requests.post(
-            f"{self.server_url}/register_peer",
-            json={"peer_id": peer_id}
-        )
-        return response.json()
+        try:
+            response = requests.post(
+                f"{self.server_url}/register_peer",
+                json={"peer_id": peer_id},
+                timeout=5  # Add timeout
+            )
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            st.error(f"Connection error: {str(e)}")
+            return {"status": "error", "message": str(e)}
     
     def send_offer(self, peer_id: str, offer: Dict) -> Dict:
-        response = requests.post(
-            f"{self.server_url}/send_offer",
-            json={"peer_id": peer_id, "offer": offer}
-        )
-        return response.json()
+        try:
+            response = requests.post(
+                f"{self.server_url}/send_offer",
+                json={"peer_id": peer_id, "offer": offer},
+                timeout=5
+            )
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            st.error(f"Connection error: {str(e)}")
+            return {"status": "error", "message": str(e)}
     
     def get_answer(self, peer_id: str) -> Dict:
-        response = requests.get(
-            f"{self.server_url}/get_answer",
-            params={"peer_id": peer_id}
-        )
-        return response.json()
+        try:
+            response = requests.get(
+                f"{self.server_url}/get_answer",
+                params={"peer_id": peer_id},
+                timeout=5
+            )
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            st.error(f"Connection error: {str(e)}")
+            return {"status": "error", "message": str(e)}
 
 def create_video_interface():
     st.set_page_config(
@@ -46,7 +62,7 @@ def create_video_interface():
 
     # Initialize session state
     if 'connection' not in st.session_state:
-        st.session_state.connection = WebRTCConnection("http://localhost:7860")
+        st.session_state.connection = WebRTCConnection()
     if 'status' not in st.session_state:
         st.session_state.status = "Initializing..."
     if 'peer_id' not in st.session_state:
@@ -83,15 +99,12 @@ def create_video_interface():
         
         if st.button("Connect", key="connect_button", use_container_width=True):
             if st.session_state.peer_id:
-                try:
-                    response = st.session_state.connection.register_peer(
-                        st.session_state.peer_id
-                    )
-                    if response.get("status") == "registered":
-                        st.session_state.status = "Connected"
-                        st.experimental_rerun()
-                except Exception as e:
-                    st.error(f"Connection failed: {str(e)}")
+                response = st.session_state.connection.register_peer(
+                    st.session_state.peer_id
+                )
+                if response.get("status") == "registered":
+                    st.session_state.status = "Connected"
+                    st.rerun()  # Using st.rerun() instead of experimental_rerun
             else:
                 st.error("Please enter or generate a Peer ID")
 
@@ -144,7 +157,7 @@ def create_video_interface():
 
     if st.button("End Call", key="end_call", use_container_width=True):
         st.session_state.status = "Call ended"
-        st.experimental_rerun()
+        st.rerun()  # Using st.rerun() instead of experimental_rerun
 
 if __name__ == "__main__":
     create_video_interface()