laurentS / slowapi

A rate limiter for Starlette and FastAPI
https://pypi.org/project/slowapi/
MIT License
1.18k stars 74 forks source link

Dynamic rate limit based on user type #13

Closed sdklab007 closed 3 years ago

sdklab007 commented 4 years ago

I need a way to dynamically set the rate limit based on the user type.

For example, I want to limit users without access token & have unlimited access to users with the access token.

What I am currently using:


limiter = Limiter(key_func=identify_my_user_func)
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

def identify_my_user_func(request: Request):
    if 'access_token' not in request.query_params:
        return request.client.host
    return "REGISTERED_USER"

@limiter.limit("2/minute")
def some_request(request: Request)):
     return data

I am trying to find a way to conditionally limit 2/minute. Basically I want to increase the limit based on the user type.

laurentS commented 4 years ago

Hi @sdklab007, you should be able to use a callable to pick the limit, like:

def get_limit_for_user():
    return "2/minute"

@limiter.limit(get_limit_for_user):
def some_request(request: Request):
    pass

and if you want some users to be exempted from the limit, you should also be able to do:

def is_user_exempt():
    pass # return a boolean

@limiter.limit(get_limit_for_user, exempt_when=is_request_exempt):
def some_request():
    pass

I hope this helps!

sdklab007 commented 4 years ago

@laurentS Thank you so much.

sdklab007 commented 4 years ago

@laurentS One more query, how do I get the request object in get_limit_for_user to identify the user.

laurentS commented 4 years ago

@sdklab007 sorry for the lag. I'm afraid I don't have a good solution for your last question. This is a use case which I don't think has been needed so far. The code was ported from flask-limiter where it's possible to access the current request object almost like a global variable (see https://github.com/encode/starlette/issues/420 for a bit more on this), and I did not think of this at the time. If you're in a hurry, you can probably hack something together based on the ticket above, but I'll add it to my todo list to change the code to handle this use case, I think the current status is not acceptable :sweat: Obviously, PRs are always welcome if you're faster to it than me! :wink:

sdklab007 commented 4 years ago

@laurentS Thanks for your kind update. I need to hack a bit as per the link you've shared.

Sure, I will see if I can contribute :)

sdklab007 commented 3 years ago

I was able to solve this by the link you had shared, below is the way if someone needs it:

REQUEST_CTX_KEY = "request_context"
_request_ctx_var: ContextVar[str] = ContextVar(REQUEST_CTX_KEY, default=None)

@app.middleware("http")
async def request_context_middleware(request: Request, call_next):
    try:
        request_ctx = _request_ctx_var.set(request)
        response = await call_next(request)
        _request_ctx_var.reset(request_ctx)
        return response
    except Exception as e:
        raise e

Cheers!! @laurentS

Bear1110 commented 2 years ago

This post help me a lot! Thanks

fredi-python commented 10 months ago

Could someone give an example, of how to use @sdklab007's code in practice?

My endpoint is the following:

@app.post("/v1/chat/completions")
@limiter.limit("2/second")
@limiter.limit("10/minute")
@limiter.limit("100/hour")
@limiter.limit("2000/day")
async def chat_completion(request: Request, data: dict = Body(...)):
    model = data.get("model", None)

I want to check if the model equals llama-70b, if, set rate limits to:

@limiter.limit("1/second")
@limiter.limit("5/minute")
@limiter.limit("50/hour")
@limiter.limit("1000/day")
gellnerm commented 9 months ago

You can also use a double-wrapper (a whopper :smiley:) to get access to the request.

def condition_func(request: Request, func, *args, **kwargs):

    if no_limit:
        return func.__wrapped__(request=request, *args, **kwargs)  # call unlimited func

    return func(request=request, *args, **kwargs)

def ratelimit(*decor_args, **decor_kwargs):

    def decorate(func):
        condition_func = decor_kwargs.pop('condition_func')
        func = decor_kwargs.pop('limiter').limit(*decor_args, **decor_kwargs)(func)

        @functools.wraps(func)
        def wrapper(request: Request, *args, **kwargs):
            return condition_func(request, func, *args, **kwargs)

        return wrapper

    return decorate

Use it like the original:

@ratelimit('10/day', limiter=limiter, condition_func=condition_func)
fredi-python commented 9 months ago

@gellnerm Could you show me a complete example in fastapi? Also on discord if you want (username: fredipy) Thanks

fredi-python commented 9 months ago

Like I want to change the rate limit based on the data that gets sent.

data = await request.json()
if data.get("model") == "llama-2-13b":
    rate_limit = "5/minute;30/hour"
seizoux commented 7 months ago

any news here? i have an async func that calls my PSQL database and returns a str containing a ratelimit, i want to use this.

npip99 commented 7 months ago

@laurentS Can we reopen this issue in the meantime? A few seem to be asking for it.

Personally I have a user_jwt = Depends(parse_jwt), where user_jwt is a JWT Token that has a user_id and a rate limit. It would be useful if slowapi has the ability to read the rate limit from their JWT Token (i.e. have access to the same parameters as the API request itself). A separate feature on a similar topic is to use their user_id as the key instead of their IP address.

It might be hard to implement, but I'm thinking of a use case like this:

@limiter.limit("4/second")
@limiter.limit("20/minute")
@limiter.limit("500/hour")
@limiter.limit("10000/day")
@limiter.limit(custom=...)
@router.post("/get-some-data")
async def get_some_data(data: dict = Body(...), user_jwt: dict = Depends(guard_login)):

Where, the first 4 use the IP address as key, and the second 4 are by user_id using the user_id as key and the JWT's rate limit. custom receives a Python function that takes in all of the parameters of the API request (data = Body(...), user_jwt = Depends(parse_jwt), db = Depends(get_db)), and returns the key and List[rate limit strings]. For @seizoux's PSQL use case, he can use db as a parameter (And the IP rate limit before it prevents abusing db connections too much).

Perhaps some or all of that is not possible, I'm not really sure what's under the hood or how this is implemented.

anu-kailash commented 3 months ago

With FastAPI becoming more and more popular, this is a common usecase now. Either request should be made accessible or async dynamic limit Callable must be supported.

Async support would probably be easier and help with most cases?

@laurentS What are your thoughts on this?

Ivareh commented 1 month ago

I would also like support for this as @anu-kailash proposes, though it may be complicated.