Repository: ragapp/ragapp
File: src/ragapp/backend/middlewares/rate_limit.py (Python)

Code before:

from backend.services.user_chat_service import UserChatService

JWT_COOKIE_NAME = "Authorization"  # The name of the cookie that stores the JWT token
JWT_USER_ID_CLAIM = "preferred_username"  # The claim in the JWT token that stores the user ID or user name
JWT_USER_ROLES_CLAIM = (
    "X-Forwarded-Roles"  # The claim in the JWT token that stores the user roles
)
CHAT_REQUEST_LIMIT_THRESHOLD = int(os.getenv("CHAT_REQUEST_LIMIT_THRESHOLD", 0))
CHAT_REQUEST_LIMIT_ENABLED = CHAT_REQUEST_LIMIT_THRESHOLD > 0
ADMIN_ROLE = "admin-manager"


class UserInfo(BaseModel):
    username: str
    roles: list[str] = []

    @computed_field
    def is_admin(self) -> bool:
        return ADMIN_ROLE in self.roles


async def request_limit_middleware(request: Request) -> Response:
    if CHAT_REQUEST_LIMIT_ENABLED:
        time_frame = _get_time_frame()
        user = _extract_user_info_from_request(request)
        # Use user name as the key for rate limiting
        request_count = UserChatService.get_user_chat_request_count(
            user.username, time_frame
        )
        if not user.is_admin and request_count >= CHAT_REQUEST_LIMIT_THRESHOLD:
            raise HTTPException(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                detail=f"You have exceeded {CHAT_REQUEST_LIMIT_THRESHOLD} chat requests. Please try again later.",
            )

        UserChatService.update_user_chat_request_count(
            user.username, time_frame, request_count + 1
        )


def _get_time_frame():
    """
    Provide the time frame for the request count
    """
    # Use current date as the time frame
    return time.strftime("%Y-%m-%d")


def _extract_user_info_from_request(request: Request) -> UserInfo:
    cookie = request.cookies.get(JWT_COOKIE_NAME)

Review comment (refactor):
good idea with `UserInfo` - but it's not a concern of the rate limiting, please move it to a own file (it's a model) and use it here like:

```suggestion
        user = UserInfo.extract_from_request(request)
```

also move the `extract_from_request` function below then

Code after:

CHAT_REQUEST_LIMIT_THRESHOLD = int(os.getenv("CHAT_REQUEST_LIMIT_THRESHOLD", 0))
CHAT_REQUEST_LIMIT_ENABLED = CHAT_REQUEST_LIMIT_THRESHOLD > 0


async def request_limit_middleware(request: Request) -> Response:
    if CHAT_REQUEST_LIMIT_ENABLED:
        user = UserInfo.from_request(request)
        time_frame = _get_time_frame()
        request_count = UserChatService.get_user_chat_request_count(user, time_frame)
        if not user.is_admin and request_count >= CHAT_REQUEST_LIMIT_THRESHOLD:
            raise HTTPException(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                detail=f"You have exceeded {CHAT_REQUEST_LIMIT_THRESHOLD} chat requests. Please try again later.",
            )

        UserChatService.update_user_chat_request_count(
            user, time_frame, request_count + 1
        )


def _get_time_frame():
    """
    Provide the time frame for the request count
    """
    # Use current date as the time frame
    return time.strftime("%Y-%m-%d")