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