chore: adding a retry.

play_with_endpoint.py

@@ -27,11 +27,17 @@ headers = {
 }
 
 
-def query(payload):
+def query(payload, allowed_retries=2):
     response = requests.post(API_URL, headers=headers, json=payload)
 
-    if "error" in response:
-        assert False, f"Got an error: {response=}"
+    if "error" in response.json():
+        if allowed_retries > 0:
+            # Sometimes we have "Bad gateway" error
+            print(f"Warning, error {response=} {response.json()=} in the query, relaunching")
+
+            return query(payload, allowed_retries - 1)
+
+        assert False, f"Got an error: {response=} {response.json()=}"
 
     return response.json()