fix(agents-api): Fix cozo query retry logic

Julep Developers · Julep Developers · commit 1cb593201e70 · 2024-11-29T14:48:26.000Z
Signed-off-by: Julep Developers &lt;developers@julep.ai&gt;
diff --git a/agents-api/agents_api/models/utils.py b/agents-api/agents_api/models/utils.py
@@ -7,7 +7,10 @@
 
 import pandas as pd
 from fastapi import HTTPException
+from httpcore import NetworkError, TimeoutException
+from httpx import RequestError
 from pydantic import BaseModel
+from requests.exceptions import ConnectionError, Timeout
 
 from ..common.utils.cozo import uuid_int_list_to_uuid4
 from ..env import do_verify_developer, do_verify_developer_owns_resource
@@ -254,7 +257,21 @@ def wrapper(*args: P.args, client=None, **kwargs: P.kwargs) -> pd.DataFrame:
 
                 debug and print(repr(e))
 
-                if "busy" in (str(e) + str(getattr(e, "resp", e))).lower():
+                pretty_error = repr(e).lower()
+                cozo_busy = ("busy" in pretty_error) or (
+                    "when executing against relation '_" in pretty_error
+                )
+                connection_error = isinstance(
+                    e,
+                    (
+                        ConnectionError,
+                        Timeout,
+                        TimeoutException,
+                        NetworkError,
+                        RequestError,
+                    ),
+                )
+                if cozo_busy or connection_error:
                     raise HTTPException(
                         status_code=429, detail="Resource busy. Please try again later."
                     ) from e
@@ -321,9 +338,10 @@ def is_resource_busy(e: Exception) -> bool:
             return isinstance(e, HTTPException) and e.status_code == 429
 
         @retry(
-            stop=stop_after_attempt(4),
-            wait=wait_exponential(multiplier=1, min=4, max=10),
+            stop=stop_after_attempt(6),
+            wait=wait_exponential(multiplier=1.2, min=3, max=10),
             retry=retry_if_exception(is_resource_busy),
+            reraise=True,
         )
         @wraps(func)
         async def wrapper(
@@ -367,7 +385,21 @@ async def wrapper(
 
                 debug and print(repr(e))
 
-                if "busy" in (str(e) + str(getattr(e, "resp", e))).lower():
+                pretty_error = repr(e).lower()
+                cozo_busy = ("busy" in pretty_error) or (
+                    "when executing against relation '_" in pretty_error
+                )
+                connection_error = isinstance(
+                    e,
+                    (
+                        ConnectionError,
+                        Timeout,
+                        TimeoutException,
+                        NetworkError,
+                        RequestError,
+                    ),
+                )
+                if cozo_busy or connection_error:
                     raise HTTPException(
                         status_code=429, detail="Resource busy. Please try again later."
                     ) from e
diff --git a/agents-api/gunicorn_conf.py b/agents-api/gunicorn_conf.py
@@ -4,7 +4,7 @@
 debug = os.getenv("AGENTS_API_DEBUG", "false").lower() == "true"
 
 # Gunicorn config variables
-workers = (multiprocessing.cpu_count() - 1) if not debug else 1
+workers = (multiprocessing.cpu_count() // 2) if not debug else 1
 worker_class = "uvicorn.workers.UvicornWorker"
 bind = "0.0.0.0:8080"
 keepalive = 120
