Update app/api/v1/endpoints/tts.py

app/api/v1/endpoints/tts.py

@@ -1,4 +1,4 @@
-from fastapi import APIRouter, HTTPException, Depends
+from fastapi import APIRouter, HTTPException, Depends, Request
 from fastapi.responses import StreamingResponse
 from sqlalchemy.orm import Session
 from ....services.tts import TTSService
@@ -6,25 +6,114 @@ from ....schemas.tts import TTSRequest, TTSResponse
 from app.models.job import ProcessingJob
 from app.schemas.job import JobCreate
 from app.core.config import get_db
-from typing import Optional
+from typing import Optional, Dict, Tuple
+import time
+from datetime import datetime, timedelta
+from fastapi import status
 
 router = APIRouter()
 
+# Rate limiting storage (use Redis in production)
+_RATE_LIMIT_STORE: Dict[str, Tuple[int, float]] = {}
+_CLEANUP_INTERVAL = 60  # Seconds between cache cleanups
+_LAST_CLEANUP = time.time()
+
+# Rate limiting configuration
+RATE_LIMIT_WINDOW = 60  # 60-second window
+MAX_REQUESTS_PER_WINDOW = 10  # Max requests per IP per window
+MIN_REQUEST_INTERVAL = 1.0  # 1 second between requests
+
+def _cleanup_rate_limit_store():
+    """Remove expired rate limit entries"""
+    global _LAST_CLEANUP
+    current_time = time.time()
+    if current_time - _LAST_CLEANUP < _CLEANUP_INTERVAL:
+        return
+    
+    expired = current_time - RATE_LIMIT_WINDOW
+    to_delete = [ip for ip, (_, timestamp) in _RATE_LIMIT_STORE.items() 
+                if timestamp < expired]
+    
+    for ip in to_delete:
+        del _RATE_LIMIT_STORE[ip]
+    
+    _LAST_CLEANUP = current_time
+
+def _check_rate_limit(client_ip: str) -> Optional[float]:
+    """Check if request is allowed, returns retry_after if rate limited"""
+    _cleanup_rate_limit_store()
+    current_time = time.time()
+    
+    # Check minimum interval
+    if client_ip in _RATE_LIMIT_STORE:
+        last_request_time = _RATE_LIMIT_STORE[client_ip][1]
+        if current_time - last_request_time < MIN_REQUEST_INTERVAL:
+            return MIN_REQUEST_INTERVAL - (current_time - last_request_time)
+    
+    # Check request count in window
+    request_count, _ = _RATE_LIMIT_STORE.get(client_ip, (0, 0))
+    if request_count >= MAX_REQUESTS_PER_WINDOW:
+        oldest_entry_time = _RATE_LIMIT_STORE[client_ip][1]
+        return RATE_LIMIT_WINDOW - (current_time - oldest_entry_time)
+    
+    return None
+
+def _update_rate_limit(client_ip: str):
+    """Update rate limit counters for the client"""
+    current_time = time.time()
+    if client_ip in _RATE_LIMIT_STORE:
+        count, _ = _RATE_LIMIT_STORE[client_ip]
+        _RATE_LIMIT_STORE[client_ip] = (count + 1, current_time)
+    else:
+        _RATE_LIMIT_STORE[client_ip] = (1, current_time)
+
 @router.post("/", response_model=TTSResponse)
 async def text_to_speech(
-    request: TTSRequest,
+    request: Request,
+    tts_request: TTSRequest,
     db: Session = Depends(get_db)
 ):
+    # Get client IP (handles proxies properly)
+    client_ip = request.client.host
+    if not client_ip:
+        client_ip = "unknown"
+    
+    # Check rate limits
+    if retry_after := _check_rate_limit(client_ip):
+        raise HTTPException(
+            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+            detail={
+                "message": "Rate limit exceeded",
+                "retry_after": retry_after,
+                "limits": {
+                    "requests_per_window": MAX_REQUESTS_PER_WINDOW,
+                    "window_seconds": RATE_LIMIT_WINDOW,
+                    "min_request_interval": MIN_REQUEST_INTERVAL
+                }
+            },
+            headers={
+                "Retry-After": str(int(retry_after)),
+                "X-RateLimit-Limit": str(MAX_REQUESTS_PER_WINDOW),
+                "X-RateLimit-Remaining": str(MAX_REQUESTS_PER_WINDOW - _RATE_LIMIT_STORE.get(client_ip, (0, 0))[0]),
+                "X-RateLimit-Reset": str(int(time.time() + retry_after))
+            }
+        )
+    
     try:
+        # Update rate limit counters
+        _update_rate_limit(client_ip)
+        
+        # Process the request
         service = TTSService()
-        audio_buffer = service.generate_audio(request.text)
+        audio_buffer = service.generate_audio(tts_request.text)
         
-        # Save to database
+        # Database logging
         job = JobCreate(
             job_type="tts",
             input_file=None,
             status="completed",
-            result=f"Generated {len(request.text)} characters"
+            result=f"Generated {len(tts_request.text)} characters",
+            client_ip=client_ip  # Track which IP made the request
         )
         
         db_job = ProcessingJob(**job.model_dump())
@@ -37,13 +126,23 @@ async def text_to_speech(
             media_type="audio/wav",
             headers={
                 "Content-Disposition": "attachment; filename=speech.wav",
-                "Text-Length": str(len(request.text)),
-                "Sample-Rate": str(service.sample_rate)
+                "Text-Length": str(len(tts_request.text)),
+                "Sample-Rate": str(service.sample_rate),
+                "X-RateLimit-Limit": str(MAX_REQUESTS_PER_WINDOW),
+                "X-RateLimit-Remaining": str(MAX_REQUESTS_PER_WINDOW - _RATE_LIMIT_STORE[client_ip][0]),
+                "X-RateLimit-Reset": str(int(_RATE_LIMIT_STORE[client_ip][1] + RATE_LIMIT_WINDOW))
             }
         )
         
-    except HTTPException:
-        raise
+    except HTTPException as e:
+        db.rollback()
+        raise e
     except Exception as e:
         db.rollback()
-        raise HTTPException(status_code=500, detail=str(e))
+        # Log failed attempts in rate limiting
+        if client_ip in _RATE_LIMIT_STORE:
+            _RATE_LIMIT_STORE[client_ip] = (_RATE_LIMIT_STORE[client_ip][0] - 1, _RATE_LIMIT_STORE[client_ip][1])
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"TTS processing failed: {str(e)}"
+        )