Fix: RateLimit requests were not released when a streaming generation…

… exception occurred (#11540)
langgenius · Dec 11, 2024 · d05f189 · d05f189
1 parent ceaa9f1
commit d05f189
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/api/core/app/features/rate_limiting/rate_limit.py b/api/core/app/features/rate_limiting/rate_limit.py
@@ -110,7 +110,7 @@ def __next__(self):
             raise StopIteration
         try:
             return next(self.generator)
-        except StopIteration:
+        except Exception:
             self.close()
             raise
 

diff --git a/api/services/app_generate_service.py b/api/services/app_generate_service.py
@@ -108,6 +108,9 @@ def generate(
                 raise ValueError(f"Invalid app mode {app_model.mode}")
         except RateLimitError as e:
             raise InvokeRateLimitError(str(e))
+        except Exception:
+            rate_limit.exit(request_id)
+            raise
         finally:
             if not streaming:
                 rate_limit.exit(request_id)