retry groq requests

GooeyAI · Jul 9, 2024 · b78304e · b78304e
1 parent cf81dce
commit b78304e
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 20 deletions.
diff --git a/daras_ai_v2/language_model.py b/daras_ai_v2/language_model.py
@@ -25,12 +25,12 @@
 from daras_ai.image_input import gs_url_to_uri, bytes_to_cv2_img, cv2_img_to_bytes
 from daras_ai_v2.asr import get_google_auth_session
 from daras_ai_v2.exceptions import raise_for_status, UserError
-from functions.recipe_functions import LLMTools
 from daras_ai_v2.gpu_server import call_celery_task
 from daras_ai_v2.text_splitter import (
     default_length_function,
     default_separators,
 )
+from functions.recipe_functions import LLMTools
 
 DEFAULT_SYSTEM_MSG = "You are an intelligent AI assistant. Follow the instructions as closely as possible."
 
@@ -985,16 +985,7 @@ def get_openai_client(model: str):
     return client
 
 
-def groq_should_retry(e: Exception) -> bool:
-    return isinstance(e, requests.HTTPError) and e.response.status_code in [
-        429,
-        502,
-        503,
-        504,
-    ]
-
-
-@aifail.retry_if(groq_should_retry)
+@aifail.retry_if(aifail.http_should_retry)
 def _run_groq_chat(
     *,
     model: str,

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -80,7 +80,7 @@ ua-parser = "^0.18.0"
 user-agents = "^2.2.0"
 openpyxl = "^3.1.2"
 loguru = "^0.7.2"
-aifail = "0.2.0"
+aifail = {git = "https://github.com/GooeyAI/aifail/", rev = "0.3.0"}
 pytest-playwright = "^0.4.3"
 emoji = "^2.10.1"
 pyvespa = "^0.39.0"