mirror of
https://github.com/khoaliber/khoj.git
synced 2026-03-02 13:18:18 +00:00
Convert websocket rate limiter to async method
Fixes calling websocket rate limiter from async chat_ws method. Not sure why the issue did not trigger in local setups. Maybe has to do with gunicorn vs uvicorn / multi-workers setup in prod vs local.
This commit is contained in:
@@ -1533,8 +1533,8 @@ async def chat_ws(
|
|||||||
|
|
||||||
# Apply rate limiting manually
|
# Apply rate limiting manually
|
||||||
try:
|
try:
|
||||||
rate_limiter_per_minute.check_websocket(websocket)
|
await rate_limiter_per_minute.check_websocket(websocket)
|
||||||
rate_limiter_per_day.check_websocket(websocket)
|
await rate_limiter_per_day.check_websocket(websocket)
|
||||||
image_rate_limiter.check_websocket(websocket, body)
|
image_rate_limiter.check_websocket(websocket, body)
|
||||||
except HTTPException as e:
|
except HTTPException as e:
|
||||||
await websocket.send_text(json.dumps({"error": e.detail}))
|
await websocket.send_text(json.dumps({"error": e.detail}))
|
||||||
|
|||||||
@@ -1936,7 +1936,7 @@ class ApiUserRateLimiter:
|
|||||||
# Add the current request to the cache
|
# Add the current request to the cache
|
||||||
UserRequests.objects.create(user=user, slug=self.slug)
|
UserRequests.objects.create(user=user, slug=self.slug)
|
||||||
|
|
||||||
def check_websocket(self, websocket: WebSocket):
|
async def check_websocket(self, websocket: WebSocket):
|
||||||
"""WebSocket-specific rate limiting method"""
|
"""WebSocket-specific rate limiting method"""
|
||||||
# Rate limiting disabled if billing is disabled
|
# Rate limiting disabled if billing is disabled
|
||||||
if state.billing_enabled is False:
|
if state.billing_enabled is False:
|
||||||
@@ -1954,7 +1954,7 @@ class ApiUserRateLimiter:
|
|||||||
|
|
||||||
# Remove requests outside of the time window
|
# Remove requests outside of the time window
|
||||||
cutoff = django_timezone.now() - timedelta(seconds=self.window)
|
cutoff = django_timezone.now() - timedelta(seconds=self.window)
|
||||||
count_requests = UserRequests.objects.filter(user=user, created_at__gte=cutoff, slug=self.slug).count()
|
count_requests = await UserRequests.objects.filter(user=user, created_at__gte=cutoff, slug=self.slug).acount()
|
||||||
|
|
||||||
# Check if the user has exceeded the rate limit
|
# Check if the user has exceeded the rate limit
|
||||||
if subscribed and count_requests >= self.subscribed_requests:
|
if subscribed and count_requests >= self.subscribed_requests:
|
||||||
@@ -1984,7 +1984,7 @@ class ApiUserRateLimiter:
|
|||||||
)
|
)
|
||||||
|
|
||||||
# Add the current request to the cache
|
# Add the current request to the cache
|
||||||
UserRequests.objects.create(user=user, slug=self.slug)
|
await UserRequests.objects.acreate(user=user, slug=self.slug)
|
||||||
|
|
||||||
|
|
||||||
class ApiImageRateLimiter:
|
class ApiImageRateLimiter:
|
||||||
|
|||||||
Reference in New Issue
Block a user