Redis Rate Limiter
This example demonstrates how to use Redis to implement a rate limiter for an Embedbase project.
Installation
Install the required dependencies in a virtual environment:
virtualenv env
source env/bin/activate
pip install embedbase redis
Start Redis as an Embedbase rate limiter
Run a Redis instance for the Embedbase rate limiter.
docker run -d --name redis -p 6379:6379 redis
Alternatively you can use Upstash (opens in a new tab) which manage Redis for you, to go to production faster.
Start Embedbase
Create a new file main.py
with the following code:
main.py
import os
from embedbase import get_app
from embedbase.database.memory_db import MemoryDatabase
from embedbase.embedding.openai import OpenAI
import redis.asyncio as redis
import uvicorn
from fastapi.responses import JSONResponse
REDIS_HOST = os.environ.get("REDIS_HOST", "localhost")
REDIS_PORT = os.environ.get("REDIS_PORT", 6379)
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
async def get_redis():
r = redis.Redis(host=REDIS_HOST, port=REDIS_PORT)
await r.ping()
return r
async def rate_limiter(request, call_next, db, embedder):
"""Rate limit middleware using Redis."""
rate_limit = 1 # Set your rate limit here (e.g., 1 requests per minute)
expiration = 60 # Set the expiration time in seconds (e.g., 60 seconds)
redis_client = await get_redis()
key = f"rate_limit:{request.client.host}"
current_requests = await redis_client.get(key)
if current_requests is None:
await redis_client.set(key, 1, ex=expiration)
elif int(current_requests) < rate_limit:
await redis_client.incr(key)
else:
return JSONResponse(
content={"detail": "Too many requests"},
status_code=429, # HTTP status code for "Too Many Requests"
)
response = await call_next(request)
return response
app = (
get_app()
.use_embedder(OpenAI(OPENAI_API_KEY))
.use_db(MemoryDatabase())
.use_middleware(rate_limiter)
.run()
)
if __name__ == "__main__":
uvicorn.run("main:app", reload=True)
Start the Embedbase application with the following command:
python3 main.py
Test Embedbase rate limiter
To test if the rate limiter is working, open another terminal and run the following commands:
Create a shell script test.sh
to send 10 parallel requests to the Embedbase API:
cat > test.sh <<EOF
for i in {1..10}; do
curl -X POST "http://localhost:8000/v1/sampledataset/search" \
-H "accept: application/json" \
-H "Content-Type: application/json" \
-d "{ \"query\": \"Hello world!\"}"
done
EOF
Run the test script:
bash test.sh
The API should now respond with a 429 "Too Many Requests" HTTP status code when the rate limit has been reached.