Update app.py

app.py

@@ -443,24 +443,26 @@ def get_instances():
 
 @app.route('/metrics/<username>/<instance_id>')
 def stream_metrics(username, instance_id):
-    # Use requests.get with stream=True for SSE
     url = f"https://api.hf.space/v1/{username}/{instance_id}/live-metrics/sse"
 
     def generate():
-      try:
-        response = requests.get(url, stream=True, headers={"Accept": "text/event-stream"})
-        response.raise_for_status()  # Important: raise HTTP errors
-        client = SSEClient(response)
-        for event in client.events():
-          if event.event == 'metric':
-            yield f"event: {event.event}\\ndata: {event.data}\\n\\n"
-      except requests.exceptions.RequestException as e:
-          print(f"Request Exception: {e}")
-          yield f"event: error\ndata: Connection error\n\n"  # Send an error event
-      except Exception as e:
-        print(f"An error occurred: {e}")
-        yield f"event: error\ndata: An error occurred\n\n"
-
+        session = requests.Session()  # Create a session object
+        try:
+            # Add a timeout (e.g., 10 seconds)
+            response = session.get(url, stream=True, headers={"Accept": "text/event-stream"}, timeout=10)
+            response.raise_for_status()  # Important: raise HTTP errors
+            client = SSEClient(response)
+            for event in client.events():
+                if event.event == 'metric':
+                    yield f"event: {event.event}\\ndata: {event.data}\\n\\n"
+        except requests.exceptions.RequestException as e:
+            print(f"Request Exception: {e}")
+            yield f"event: error\ndata: Connection error: {e}\\n\\n"  # Include error details
+        except Exception as e:
+            print(f"An error occurred: {e}")
+            yield f"event: error\ndata: An error occurred: {e}\\n\\n"  # Include error details
+        finally:
+            session.close() # Close the session
 
     return Response(generate(), mimetype='text/event-stream')