Repository: LAION-AI/Open-Assistant
File: backend/oasst_backend/api/deps.py (Python)

Code before:
        # (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 UserTaskTypeRateLimiter(RateLimiter):
    """
    User-level rate limiter for a specific task type.
    """

    def __init__(
        self,
        task_types: list[str],
        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)
        self.task_types = task_types

    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
        json = await request.json()
        user = 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

        # Skip when the request is not in our task types of interest
        if not json.get("type") or json.get("type") not in self.task_types:
            return

        return await super().__call__(request, response)



Review comment (security):
To be guaranteed unique the key needs to include api-key, auth-method and user-id, beside of extracting the user from the request json it would probably be easier to get it from the oasst_user_header (`x-oasst-user`) .. which already contains auth-method+username delimited by ':'.

Code after:
            return

        # Attempt to retrieve api_key and user information
        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 UserTaskTypeRateLimiter(RateLimiter):
    """
    User-level rate limiter for a specific task type.
    """

    def __init__(
        self,
        task_types: list[str],
        times: int = 100,
        milliseconds: int = 0,
        seconds: int = 0,
        minutes: int = 1,
        hours: int = 0,
    ) -> None:
        super().__init__(times, milliseconds, seconds, minutes, hours, user_identifier)
        self.task_types = task_types

    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
        json = await request.json()
        user = 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

        # Skip when the request is not in our task types of interest
        if not json.get("type") or json.get("type") not in self.task_types:
            return

        return await super().__call__(request, response)