From 44f086b7da5a5bce815eb1adb90d18b20f87670c Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Tue, 28 May 2024 10:07:45 -0700 Subject: [PATCH] fix: rate limit filter --- pipelines/examples/rate_limit_filter_pipeline.py | 13 +++++++++---- pipelines/rate_limit_filter_pipeline.py | 13 +++++++++---- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/pipelines/examples/rate_limit_filter_pipeline.py b/pipelines/examples/rate_limit_filter_pipeline.py index a3bea03..4bd9112 100644 --- a/pipelines/examples/rate_limit_filter_pipeline.py +++ b/pipelines/examples/rate_limit_filter_pipeline.py @@ -20,10 +20,8 @@ class Pipeline: # List target pipelines (models) that this filter will be connected to. self.pipelines = ["*"] - pass - # Initialize rate limits - self.requests_per_minute: Optional[int] = 60 + self.requests_per_minute: Optional[int] = 10 self.requests_per_hour: Optional[int] = 1000 self.sliding_window_limit: Optional[int] = 100 self.sliding_window_minutes: Optional[int] = 15 @@ -48,7 +46,14 @@ class Pipeline: self.user_requests[user_id] = [ req for req in self.user_requests[user_id] - if now - req < self.sliding_window_minutes * 60 + if ( + (self.requests_per_minute is not None and now - req < 60) + or (self.requests_per_hour is not None and now - req < 3600) + or ( + self.sliding_window_limit is not None + and now - req < self.sliding_window_minutes * 60 + ) + ) ] def log_request(self, user_id: str): diff --git a/pipelines/rate_limit_filter_pipeline.py b/pipelines/rate_limit_filter_pipeline.py index a3bea03..4bd9112 100644 --- a/pipelines/rate_limit_filter_pipeline.py +++ b/pipelines/rate_limit_filter_pipeline.py @@ -20,10 +20,8 @@ class Pipeline: # List target pipelines (models) that this filter will be connected to. self.pipelines = ["*"] - pass - # Initialize rate limits - self.requests_per_minute: Optional[int] = 60 + self.requests_per_minute: Optional[int] = 10 self.requests_per_hour: Optional[int] = 1000 self.sliding_window_limit: Optional[int] = 100 self.sliding_window_minutes: Optional[int] = 15 @@ -48,7 +46,14 @@ class Pipeline: self.user_requests[user_id] = [ req for req in self.user_requests[user_id] - if now - req < self.sliding_window_minutes * 60 + if ( + (self.requests_per_minute is not None and now - req < 60) + or (self.requests_per_hour is not None and now - req < 3600) + or ( + self.sliding_window_limit is not None + and now - req < self.sliding_window_minutes * 60 + ) + ) ] def log_request(self, user_id: str):