Spaces:
Sleeping
Sleeping
File size: 3,289 Bytes
2a8faae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
"""
Middleware for Medical RAG AI Advisor API
"""
import time
import logging
from typing import Callable, Awaitable
from fastapi import Request, Response, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware.base import BaseHTTPMiddleware
logger = logging.getLogger(__name__)
class ProcessTimeMiddleware(BaseHTTPMiddleware):
"""Middleware to add processing time to response headers"""
async def dispatch(self, request: Request, call_next: Callable) -> Response:
start_time = time.time()
response = await call_next(request)
process_time = time.time() - start_time
response.headers["X-Process-Time"] = f"{process_time:.4f}"
return response
class LoggingMiddleware(BaseHTTPMiddleware):
"""Middleware for request/response logging"""
async def dispatch(self, request: Request, call_next: Callable) -> Response:
start_time = time.time()
# Log request
logger.info(f"Request: {request.method} {request.url}")
try:
response = await call_next(request)
process_time = time.time() - start_time
# Log response
logger.info(
f"Response: {response.status_code} - "
f"Time: {process_time:.4f}s - "
f"Path: {request.url.path}"
)
return response
except Exception as e:
process_time = time.time() - start_time
logger.error(
f"Error: {str(e)} - "
f"Time: {process_time:.4f}s - "
f"Path: {request.url.path}"
)
raise
class RateLimitMiddleware(BaseHTTPMiddleware):
"""Simple rate limiting middleware"""
def __init__(self, app, calls_per_minute: int = 60):
super().__init__(app)
self.calls_per_minute = calls_per_minute
self.client_calls = {}
async def dispatch(self, request: Request, call_next: Callable) -> Response:
client_ip = request.client.host
current_time = time.time()
# Clean old entries
self.client_calls = {
ip: calls for ip, calls in self.client_calls.items()
if any(call_time > current_time - 60 for call_time in calls)
}
# Check rate limit
if client_ip in self.client_calls:
recent_calls = [
call_time for call_time in self.client_calls[client_ip]
if call_time > current_time - 60
]
if len(recent_calls) >= self.calls_per_minute:
raise HTTPException(
status_code=429,
detail="Rate limit exceeded. Please try again later."
)
self.client_calls[client_ip] = recent_calls + [current_time]
else:
self.client_calls[client_ip] = [current_time]
return await call_next(request)
def get_cors_middleware_config():
"""Get CORS middleware configuration"""
return {
"allow_origins": ["*"], # Configure appropriately for production
"allow_credentials": True,
"allow_methods": ["*"],
"allow_headers": ["*"],
}
|