Add rate limiter for task request

backend/oasst_backend/api/deps.py

@@ -3,8 +3,9 @@ from secrets import token_hex
 from typing import Generator
 from uuid import UUID
 
-from fastapi import Depends, Security
+from fastapi import Depends, Request, Response, Security
 from fastapi.security.api_key import APIKey, APIKeyHeader, APIKeyQuery
+from fastapi_limiter.depends import RateLimiter
 from loguru import logger
 from oasst_backend.config import settings
 from oasst_backend.database import engine
@@ -84,3 +85,60 @@ def get_trusted_api_client(
             http_status_code=HTTPStatus.FORBIDDEN,
         )
     return client
+
+
+class UserRateLimiter(RateLimiter):
+    def __init__(
+        self, times: int = 100, milliseconds: int = 0, seconds: int = 0, minutes: int = 1, hours: int = 0
+    ) -> None:
+        async def identifier(request: Request) -> str:
+            """Identify a request based on api_key and user.id"""
+            api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key")
+            user = (await request.json()).get("user")
+            return f"{api_key}:{user.get('id')}"
+
+        super().__init__(times, milliseconds, seconds, minutes, hours, identifier)
+
+    async def __call__(self, request: Request, response: Response, api_key: str = Depends(get_api_key)) -> None:
+        # Skip if rate limiting is disabled
+        if not settings.RATE_LIMIT:
+            return
+
+        # Attempt to retrieve api_key and user information
+        api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key")
+        user = (await request.json()).get("user")
+
+        # Skip when api_key and user information are not available
+        # (such that it will be handled by `APIClientRateLimiter`)
+        if not api_key or not user or not user.get("id"):
+            return
+
+        return await super().__call__(request, response)
+
+
+class APIClientRateLimiter(RateLimiter):
+    def __init__(
+        self, times: int = 10_000, milliseconds: int = 0, seconds: int = 0, minutes: int = 1, hours: int = 0
+    ) -> None:
+        async def identifier(request: Request) -> str:
+            """Identify a request based on api_key and user.id"""
+            api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key")
+            return f"{api_key}"
+
+        super().__init__(times, milliseconds, seconds, minutes, hours, identifier)
+
+    async def __call__(self, request: Request, response: Response, api_key: str = Depends(get_api_key)) -> None:
+        # Skip if rate limiting is disabled
+        if not settings.RATE_LIMIT:
+            return
+
+        # Attempt to retrieve api_key and user information
+        api_key = request.headers.get("X-API-Key") or request.query_params.get("api_key")
+        user = (await request.json()).get("user")
+
+        # Skip if user information is available
+        # (such that it will be handled by `UserRateLimiter`)
+        if not api_key or user:
+            return
+
+        return await super().__call__(request, response)

backend/oasst_backend/api/v1/tasks.py

@@ -127,7 +127,14 @@ def generate_task(
     return task, message_tree_id, parent_message_id
 
 
-@router.post("/", response_model=protocol_schema.AnyTask)  # work with Union once more types are added
+@router.post(
+    "/",
+    response_model=protocol_schema.AnyTask,
+    dependencies=[
+        Depends(deps.UserRateLimiter(times=100, minutes=5)),
+        Depends(deps.APIClientRateLimiter(times=10_000, minutes=1)),
+    ],
+)  # work with Union once more types are added
 def request_task(
     *,
     db: Session = Depends(deps.get_db),