🐛 fix(sdk): Retry request on 500 error in upscale.py and generate_ima…

…ge/__init__.py

playground/generate_image.py

@@ -38,6 +38,7 @@ async def main():
             session=globe_s, remove_sign=True
         )
     except APIError as e:
+        print(str(e))
         print(e.response)
         return
 

pyproject.toml

@@ -20,6 +20,7 @@ dependencies = [
     "argon2-cffi>=23.1.0",
     "opencv-python>=4.8.1.78",
     "fake-useragent>=1.4.0",
+    "tenacity>=8.2.3",
 ]
 requires-python = ">=3.8"
 readme = "README.md"

src/novelai_python/sdk/ai/generate_image/__init__.py

@@ -18,6 +18,7 @@
 from curl_cffi.requests import AsyncSession
 from loguru import logger
 from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator, model_validator, Field
+from tenacity import retry, stop_after_attempt, wait_random, retry_if_exception
 from typing_extensions import override
 
 from ._enum import Model, Sampler, NoiseSchedule, ControlNetModel, Action, UCPreset
@@ -327,6 +328,12 @@ async def necessary_headers(self, request_data) -> dict:
             "Cache-Control": "no-cache",
         }
 
+    @retry(
+        wait=wait_random(min=1, max=3),
+        stop=stop_after_attempt(3),
+        retry=retry_if_exception(lambda e: hasattr(e, "code") and str(e.code) == "500"),
+        reraise=True
+    )
     async def request(self,
                       session: Union[AsyncSession, "CredentialBase"],
                       *,

src/novelai_python/sdk/ai/upscale.py

@@ -15,6 +15,7 @@
 from curl_cffi.requests import AsyncSession
 from loguru import logger
 from pydantic import ConfigDict, PrivateAttr, model_validator
+from tenacity import wait_random, retry, stop_after_attempt, retry_if_exception
 
 from ..schema import ApiBaseModel
 from ..._exceptions import APIError, AuthError, SessionHttpError
@@ -87,6 +88,12 @@ async def necessary_headers(self, request_data) -> dict:
             "Cache-Control": "no-cache",
         }
 
+    @retry(
+        wait=wait_random(min=1, max=3),
+        stop=stop_after_attempt(3),
+        retry=retry_if_exception(lambda e: hasattr(e, "code") and str(e.code) == "500"),
+        reraise=True
+    )
     async def request(self,
                       session: Union[AsyncSession, "CredentialBase"],
                       *,