Update app.py

app.py

@@ -48,7 +48,7 @@ def initialize_agent(endpoint_uri: str):
     ) 
     
 
-def wake_up_endpoint(endpoint_uri, max_wait=30):
+def wake_up_endpoint(endpoint_uri, ui, max_wait=30):
     """Poll the endpoint until it responds OK or timeout."""
 
     # Get token from environment variable or secrets store
@@ -60,22 +60,22 @@ def wake_up_endpoint(endpoint_uri, max_wait=30):
 
     while time.time() - start < max_wait:
         try:
-            print("Pinging endpoint...")
-            print(endpoint_uri)
-            response = requests.post(endpoint_uri, json=payload, timeout=5)
+            ui.append("Pinging endpoint...")
+            ui.append(endpoint_uri)
+            response = requests.post(endpoint_uri, headers=headers, json=payload, timeout=5)
             
             if response.ok:
-                print("✅ Endpoint is awake.")
+                ui.append("✅ Endpoint is awake.")
                 return True
 
             # If it's a 503 or 504, it's likely still warming up
-            print(f"Response: {response.status_code}, retrying...")
+            ui.append(f"Response: {response.status_code}, retrying...")
         except requests.exceptions.RequestException as e:
-            print(f"Exception: {e}, retrying...")
+            ui.append(f"Exception: {e}, retrying...")
 
         time.sleep(3)  # fixed short backoff; can be exponential if needed
 
-    print("❌ Timed out waiting for endpoint to wake up.")
+    ui.append("❌ Timed out waiting for endpoint to wake up.")
     return False
 
 def main():
@@ -99,7 +99,7 @@ def main():
     ui.append(endpoint_uri)   
 
     # Wake it up before health check
-    wake_up_successful = wake_up_endpoint(endpoint_uri)
+    wake_up_successful = wake_up_endpoint(endpoint_uri, ui)
     
     if not wake_up_successful:
         ui.append("Warning: Could not wake up the endpoint.  Continuing.")