Merge pull request #24 from upstash/chat-ratelimit-error

Raise a specific error for chat rate limits
upstash · Jul 2, 2024 · 43f92d9 · 43f92d9
2 parents 5fcefe9 + 646c0da
commit 43f92d9
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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}"