laurentS/slowapi

Dynamic rate limit based on user type

sdklab007 opened this issue ยท 8 comments

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.

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!

@laurentS Thank you so much.

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

@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 encode/starlette#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 ๐Ÿ˜“
Obviously, PRs are always welcome if you're faster to it than me! ๐Ÿ˜‰

@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 :)

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

This post help me a lot! Thanks

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")