Spaces:
Sleeping
Sleeping
File size: 2,235 Bytes
d2237d8 |
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 |
from fastapi import FastAPI, Depends, HTTPException, Request
from fastapi.security import APIKeyHeader
from pydantic import BaseModel
from model import ner_pipeline
import logging
import time
import json
import os
import secrets
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = FastAPI()
limiter = Limiter(key_func=get_remote_address)
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
API_KEY_FILE = "api_keys.json"
token_header = APIKeyHeader(name="X-API-KEY")
if os.path.exists(API_KEY_FILE):
with open(API_KEY_FILE, "r") as f:
API_KEYS_STORE = json.load(f)
else:
API_KEYS_STORE = {"users": {}}
API_KEYS = API_KEYS_STORE.get("users", {})
ADMIN_KEY = os.getenv("ADMIN_KEY")
class TextRequest(BaseModel):
text: str
class RegisterRequest(BaseModel):
label: str = "user"
def save_keys():
API_KEYS_STORE["users"] = API_KEYS
with open(API_KEY_FILE, "w") as f:
json.dump(API_KEYS_STORE, f, indent=2)
def verify_token(x_api_key: str = Depends(token_header)):
if x_api_key not in API_KEYS:
raise HTTPException(status_code=403, detail="Unauthorized")
return x_api_key
@app.post("/ner")
@limiter.limit("10/minute")
def ner_predict(request: TextRequest, api_key: str = Depends(verify_token), req: Request = None):
logger.info("Received NER request from IP: %s", get_remote_address(req))
predictions = ner_pipeline(request.text)
API_KEYS[api_key]["usage_count"] = API_KEYS[api_key].get("usage_count", 0) + 1
save_keys()
return {
"entities": predictions,
"usage": API_KEYS[api_key]["usage_count"]
}
@app.post("/register_user")
def register_user(request: RegisterRequest, x_api_key: str = Depends(token_header)):
if x_api_key != ADMIN_KEY:
raise HTTPException(status_code=403, detail="Admin access required")
new_key = secrets.token_urlsafe(32)
API_KEYS[new_key] = {"usage_count": 0, "label": request.label}
save_keys()
return {"message": "User registered", "api_key": new_key}
|