Rate limiting

app.py

@@ -3,7 +3,7 @@ import json
 import time
 import numpy as np
 from sentence_transformers import SentenceTransformer
-from fastapi import FastAPI, HTTPException, BackgroundTasks
+from fastapi import FastAPI, HTTPException, BackgroundTasks, Request
 from fastapi.responses import StreamingResponse, Response
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
@@ -18,6 +18,9 @@ from google import genai
 from google.genai import types
 import httpx
 from elevenlabs import ElevenLabs, VoiceSettings
+from slowapi import Limiter, _rate_limit_exceeded_handler
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -25,6 +28,18 @@ logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
+# Initialize rate limiter
+limiter = Limiter(key_func=get_remote_address)
+app.state.limiter = limiter
+
+# Custom rate limit exceeded handler with logging
+async def custom_rate_limit_handler(request: Request, exc: RateLimitExceeded):
+    client_ip = get_remote_address(request)
+    logger.warning(f"Rate limit exceeded for IP {client_ip} on endpoint {request.url.path}")
+    return await _rate_limit_exceeded_handler(request, exc)
+
+app.add_exception_handler(RateLimitExceeded, custom_rate_limit_handler)
+
 # Add CORS middleware to handle cross-origin requests
 app.add_middleware(
     CORSMiddleware,
@@ -331,20 +346,22 @@ def get_ram_usage():
     }
 
 @app.post("/api/predict")
-async def predict(request: QueryRequest):
-    query = request.query
-    history = request.history
+@limiter.limit("5/minute")  # Allow 10 chat requests per minute per IP
+async def predict(request: Request, query_request: QueryRequest):
+    query = query_request.query
+    history = query_request.history
     return StreamingResponse(stream_response(query, history), media_type="text/event-stream")
 
 @app.post("/api/tts")
-async def text_to_speech(request: TTSRequest):
+@limiter.limit("5/minute")  # Allow 5 TTS requests per minute per IP (more restrictive as TTS is more expensive)
+async def text_to_speech(request: Request, tts_request: TTSRequest):
     """Convert text to speech using ElevenLabs API"""
     if not elevenlabs_client:
         raise HTTPException(status_code=503, detail="TTS service not available")
     
     try:
         # Clean the text for TTS (remove markdown and special characters)
-        clean_text = request.text.replace("**", "").replace("*", "").replace("\n", " ").strip()
+        clean_text = tts_request.text.replace("**", "").replace("*", "").replace("\n", " ").strip()
         
         if not clean_text:
             raise HTTPException(status_code=400, detail="No text provided for TTS")
@@ -381,11 +398,13 @@ async def text_to_speech(request: TTSRequest):
         raise HTTPException(status_code=500, detail=f"TTS conversion failed: {str(e)}")
 
 @app.get("/health")
-async def health_check():
+@limiter.limit("30/minute")  # Allow frequent health checks
+async def health_check(request: Request):
     return {"status": "healthy"}
 
 @app.get("/model_info")
-async def model_info():
+@limiter.limit("10/minute")  # Limit model info requests
+async def model_info(request: Request):
     base_info = {
         "embedding_model": sentence_transformer_model,
         "faiss_index_size": len(cv_chunks),
@@ -411,7 +430,8 @@ async def model_info():
     return base_info
 
 @app.get("/ram_usage")
-async def ram_usage():
+@limiter.limit("20/minute")  # Allow moderate monitoring requests
+async def ram_usage(request: Request):
     """Endpoint to get current RAM usage."""
     try:
         ram_stats = get_ram_usage()

requirements.txt

@@ -10,4 +10,5 @@ google-genai
 asyncio
 elevenlabs
 httpx
-llama-cpp-python==0.2.85
+llama-cpp-python==0.2.85
+slowapi==0.1.9