# main.py
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from Token_bucket import TokenBucket

app = FastAPI()

class RateLimiterMiddleware(BaseHTTPMiddleware):
    def __init__(self, app, bucket: TokenBucket):
        super().__init__(app)
        self.bucket = bucket

    async def dispatch(self, request: Request, call_next):
        if self.bucket.take_token():
            return await call_next(request)

        # Return a JSON response for rate limit exceeded
        return JSONResponse(status_code=429, content={"detail": "Rate limit exceeded"})

# Token bucket with a capacity of 3 and refill rate of 1 token per second
bucket = TokenBucket(capacity=3, refill_rate=3)

# Apply the rate limiter middleware
app.add_middleware(RateLimiterMiddleware, bucket=bucket)

@app.get("/")
async def read_root():
    return {"message": "Hello World"}

# Token_bucket.py
import time

class TokenBucket:
    def __init__(self, capacity, refill_rate):
        self.capacity = capacity
        self.refill_rate = refill_rate
        self.tokens = capacity
        self.last_refill = time.time()

    def add_tokens(self):
        now = time.time()
        if self.tokens < self.capacity:
            tokens_to_add = (now - self.last_refill) * self.refill_rate
            self.tokens = min (self.capacity, self.tokens + tokens_to_add)
        self.last_refill=now

    def take_token(self):
        self.add_tokens()
        if self.tokens >= 1:
            self.tokens -=1
            return True
        return False

#!/usr/bin/env bash

ENDPOINT=${1:-"/"}  # Default endpoint to "/"
SLEEP=${2:-"0.1"}     # Default sleep time to "0.1"
COUNT=${3:-100}      # Default count to "100"

# Print the parameters for clarity
echo "Sending requests to ${ENDPOINT} with ${SLEEP}s interval, ${COUNT} times..."

# Send a request in $SLEEP seconds with a loop
for (( i=1; i<=COUNT; i++ )); do
    curl -H "Service-Name: service-1" http://127.0.0.1:8000${ENDPOINT}
    echo 
    sleep $SLEEP
done

poetry run fastapi run main.py

user@host fastapi_ratelimit % poetry run fastapi run main.py
...
INFO:     Started server process [12429]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)
INFO:     127.0.0.1:64242 - "GET / HTTP/1.1" 200 OK
INFO:     127.0.0.1:64243 - "GET / HTTP/1.1" 200 OK
INFO:     127.0.0.1:64244 - "GET / HTTP/1.1" 200 OK
INFO:     127.0.0.1:64245 - "GET / HTTP/1.1" 200 OK
INFO:     127.0.0.1:64251 - "GET / HTTP/1.1" 429 Too Many Requests
INFO:     127.0.0.1:64252 - "GET / HTTP/1.1" 429 Too Many Requests
INFO:     127.0.0.1:64253 - "GET / HTTP/1.1" 200 OK
INFO:     127.0.0.1:64254 - "GET / HTTP/1.1" 429 Too Many Requests
INFO:     127.0.0.1:64255 - "GET / HTTP/1.1" 200 OK
INFO:     127.0.0.1:64256 - "GET / HTTP/1.1" 429 Too Many Requests
INFO:     127.0.0.1:64257 - "GET / HTTP/1.1" 429 Too Many Requests
INFO:     127.0.0.1:64258 - "GET / HTTP/1.1" 200 OK
INFO:     127.0.0.1:64264 - "GET / HTTP/1.1" 429 Too Many Requests

user@host fastapi_ratelimit % ./test.sh 
Sending requests to / with 0.1s interval, 100 times...
{"message":"Hello World"}
{"message":"Hello World"}
{"message":"Hello World"}
{"message":"Hello World"}
{"detail":"Rate limit exceeded"}
{"detail":"Rate limit exceeded"}
{"message":"Hello World"}
{"detail":"Rate limit exceeded"}
{"message":"Hello World"}
{"detail":"Rate limit exceeded"}
{"detail":"Rate limit exceeded"}
{"message":"Hello World"}
{"detail":"Rate limit exceeded"}
{"detail":"Rate limit exceeded"}
{"message":"Hello World"}
{"detail":"Rate limit exceeded"}
...