From 646c0dae9957ba29d03d854210886e3f744535dc Mon Sep 17 00:00:00 2001 From: Metin Dumandag <29387993+mdumandag@users.noreply.github.com> Date: Tue, 2 Jul 2024 10:14:23 +0300 Subject: [PATCH] Raise a specific error for chat rate limits This part was not ported when rewriting the client, so I added it back. Also, updated the ruff dev dependency to the latest version. --- pyproject.toml | 2 +- upstash_qstash/errors.py | 23 +++++++++++++++++++++++ upstash_qstash/http.py | 35 ++++++++++++++++++++++++++++++----- 3 files changed, 54 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3de1ec1..bfee058 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,7 @@ pytest = "^8.2.2" python-dotenv = "^1.0.1" pytest-asyncio = "^0.23.7" mypy = "^1.10.0" -ruff = "^0.4.8" +ruff = "^0.5.0" [build-system] requires = ["poetry-core"] diff --git a/upstash_qstash/errors.py b/upstash_qstash/errors.py index fd41cb9..7e583bf 100644 --- a/upstash_qstash/errors.py +++ b/upstash_qstash/errors.py @@ -19,3 +19,26 @@ def __init__( self.limit = limit self.remaining = remaining self.reset = reset + + +class ChatRateLimitExceededError(QStashError): + def __init__( + self, + limit_requests: Optional[str], + limit_tokens: Optional[str], + remaining_requests: Optional[str], + remaining_tokens: Optional[str], + reset_requests: Optional[str], + reset_tokens: Optional[str], + ): + super( + f"Exceeded chat rate limit: " + f"Request limit: {limit_requests}, remaining: {remaining_requests}, reset: {reset_requests}; " + f"token limit: {limit_tokens}, remaining: {remaining_tokens}, reset: {reset_tokens}" + ) + self.limit_requests = limit_requests + self.limit_tokens = limit_tokens + self.remaining_requests = remaining_requests + self.remaining_tokens = remaining_tokens + self.reset_requests = reset_requests + self.reset_tokens = reset_tokens diff --git a/upstash_qstash/http.py b/upstash_qstash/http.py index b3d9933..73088e1 100644 --- a/upstash_qstash/http.py +++ b/upstash_qstash/http.py @@ -4,7 +4,11 @@ import httpx -from upstash_qstash.errors import RateLimitExceededError, QStashError +from upstash_qstash.errors import ( + RateLimitExceededError, + QStashError, + ChatRateLimitExceededError, +) class RetryConfig(TypedDict, total=False): @@ -40,10 +44,31 @@ def raise_for_non_ok_status(response: httpx.Response) -> None: return if response.status_code == 429: - limit = response.headers.get("Burst-RateLimit-Limit") - remaining = response.headers.get("Burst-RateLimit-Remaining") - reset = response.headers.get("Burst-RateLimit-Reset") - raise RateLimitExceededError(limit, remaining, reset) + headers = response.headers + if "x-ratelimit-limit-requests" in headers: + limit_requests = headers.get("x-ratelimit-limit-requests") + limit_tokens = headers.get("x-ratelimit-limit-tokens") + remaining_requests = headers.get("x-ratelimit-remaining-requests") + remaining_tokens = headers.get("x-ratelimit-remaining-tokens") + reset_requests = headers.get("x-ratelimit-reset-requests") + reset_tokens = headers.get("x-ratelimit-reset-tokens") + raise ChatRateLimitExceededError( + limit_requests=limit_requests, + limit_tokens=limit_tokens, + remaining_requests=remaining_requests, + remaining_tokens=remaining_tokens, + reset_requests=reset_requests, + reset_tokens=reset_tokens, + ) + + limit = headers.get("Burst-RateLimit-Limit") + remaining = headers.get("Burst-RateLimit-Remaining") + reset = headers.get("Burst-RateLimit-Reset") + raise RateLimitExceededError( + limit=limit, + remaining=remaining, + reset=reset, + ) raise QStashError( f"Request failed with status: {response.status_code}, body: {response.text}"