Merge pull request #26 from upstash/daily-rate-limit-error

Handle daily rate limit error separately
upstash · Jul 10, 2024 · bc329dd · bc329dd
2 parents 5be3f32 + 7c5351f
commit bc329dd
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 24 deletions.
diff --git a/upstash_qstash/errors.py b/upstash_qstash/errors.py
@@ -14,7 +14,19 @@ def __init__(
         self, limit: Optional[str], remaining: Optional[str], reset: Optional[str]
     ):
         super().__init__(
-            f"Exceeded rate limit: Limit: {limit}, remaining: {remaining}, reset: {reset}"
+            f"Exceeded burst rate limit: Limit: {limit}, remaining: {remaining}, reset: {reset}"
+        )
+        self.limit = limit
+        self.remaining = remaining
+        self.reset = reset
+
+
+class DailyMessageLimitExceededError(QStashError):
+    def __init__(
+        self, limit: Optional[str], remaining: Optional[str], reset: Optional[str]
+    ):
+        super().__init__(
+            f"Exceeded daily message limit: Limit: {limit}, remaining: {remaining}, reset: {reset}"
         )
         self.limit = limit
         self.remaining = remaining

diff --git a/upstash_qstash/http.py b/upstash_qstash/http.py
@@ -8,6 +8,7 @@
     RateLimitExceededError,
     QStashError,
     ChatRateLimitExceededError,
+    DailyMessageLimitExceededError,
 )
 
 
@@ -39,36 +40,57 @@ class RetryConfig(TypedDict, total=False):
 HttpMethod = Literal["GET", "POST", "PUT", "DELETE", "PATCH"]
 
 
+def chat_rate_limit_error(headers: httpx.Headers) -> ChatRateLimitExceededError:
+    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")
+    return 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,
+    )
+
+
+def daily_message_limit_error(headers: httpx.Headers) -> DailyMessageLimitExceededError:
+    limit = headers.get("RateLimit-Limit")
+    remaining = headers.get("RateLimit-Remaining")
+    reset = headers.get("RateLimit-Reset")
+    return DailyMessageLimitExceededError(
+        limit=limit,
+        remaining=remaining,
+        reset=reset,
+    )
+
+
+def burst_rate_limit_error(headers: httpx.Headers) -> RateLimitExceededError:
+    limit = headers.get("Burst-RateLimit-Limit")
+    remaining = headers.get("Burst-RateLimit-Remaining")
+    reset = headers.get("Burst-RateLimit-Reset")
+    return RateLimitExceededError(
+        limit=limit,
+        remaining=remaining,
+        reset=reset,
+    )
+
+
 def raise_for_non_ok_status(response: httpx.Response) -> None:
     if response.is_success:
         return
 
     if response.status_code == 429:
         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 chat_rate_limit_error(headers)
+        elif "RateLimit-Limit" in headers:
+            raise daily_message_limit_error(headers)
+        else:
+            raise burst_rate_limit_error(headers)
 
     raise QStashError(
         f"Request failed with status: {response.status_code}, body: {response.text}"