From b78304e7e8cf6d5112e39d5824d0c34ed205d99d Mon Sep 17 00:00:00 2001
From: Dev Aggarwal <devxpy@gmail.com>
Date: Tue, 9 Jul 2024 07:01:29 +0530
Subject: [PATCH] retry groq requests

---
 daras_ai_v2/language_model.py | 13 ++-----------
 poetry.lock                   | 20 ++++++++++++--------
 pyproject.toml                |  2 +-
 3 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/daras_ai_v2/language_model.py b/daras_ai_v2/language_model.py
index 8c601f2fc..463c5ecae 100644
--- 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
index 70ce360c0..c63dd4849 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -13,17 +13,21 @@ files = [
 
 [[package]]
 name = "aifail"
-version = "0.2.0"
+version = "0.3.0"
 description = ""
 optional = false
-python-versions = ">=3.10,<4.0"
-files = [
-    {file = "aifail-0.2.0-py3-none-any.whl", hash = "sha256:83f3a842dbe523ee10a4d53ee00f06e794176122b93b57752566fb98d60db603"},
-    {file = "aifail-0.2.0.tar.gz", hash = "sha256:d3e19e16740577181922055883a00e894d11413d12803e8d443094846040d6df"},
-]
+python-versions = "^3.10"
+files = []
+develop = false
 
 [package.dependencies]
-loguru = ">=0.7.2,<0.8.0"
+loguru = "^0.7.2"
+
+[package.source]
+type = "git"
+url = "https://github.com/GooeyAI/aifail/"
+reference = "0.3.0"
+resolved_reference = "577487ad5ccbf3d70d4610a6b7d09e45a4951a73"
 
 [[package]]
 name = "aiohttp"
@@ -6438,4 +6442,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.10,<3.13"
-content-hash = "fa29696b70dce0df83af18fb12a0f311dd05dccabfc7d3fe9ffe71630e042f14"
+content-hash = "cc79f3b414323945ade371a12c4071eb50b9988715f0d094e4e9ef34008c3fe2"
diff --git a/pyproject.toml b/pyproject.toml
index d703612fc..8f81fb681 100644
--- 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"