4b766cb908
- 4 provider adapters: OpenAI (SDK), Anthropic (SDK), Google (google-genai), Mistral (direct HTTP) - Core router with automatic failover + exponential backoff - Flask blueprint with /api/v1/ai/* endpoints - Auth via token-broker verify endpoint - DB models for ai_providers, ai_model_mapping, ai_router_log - /health endpoint (parallel provider check), /usage stats - 21 unit tests (all passing)
94 lines
2.8 KiB
Python
94 lines
2.8 KiB
Python
import logging
|
|
import os
|
|
import sys
|
|
from functools import wraps
|
|
|
|
from flask import request, jsonify
|
|
|
|
logger = logging.getLogger("ai_router")
|
|
|
|
TOKEN_BROKER_URL = os.environ.get(
|
|
"TOKEN_BROKER_URL", "http://localhost:8783"
|
|
)
|
|
|
|
|
|
def verify_token_via_broker(token: str) -> dict:
|
|
"""Verify an API token via the token-broker /verify endpoint."""
|
|
import requests
|
|
try:
|
|
resp = requests.post(
|
|
f"{TOKEN_BROKER_URL}/api/v1/tokens/verify",
|
|
json={"token": token},
|
|
timeout=10,
|
|
)
|
|
if resp.status_code == 200:
|
|
data = resp.json()
|
|
if data.get("valid"):
|
|
return data
|
|
return {}
|
|
except requests.RequestException as e:
|
|
logger.warning(f"Token broker unreachable: {e}")
|
|
return {}
|
|
|
|
|
|
def require_auth(f):
|
|
"""Decorator: validate Bearer or X-API-Key via token-broker."""
|
|
@wraps(f)
|
|
def decorated(*args, **kwargs):
|
|
auth_header = request.headers.get("Authorization", "")
|
|
api_key = request.headers.get("X-API-Key", "")
|
|
|
|
raw_token = ""
|
|
if auth_header.startswith("Bearer "):
|
|
raw_token = auth_header.split(" ", 1)[1]
|
|
elif api_key:
|
|
raw_token = api_key
|
|
|
|
if not raw_token:
|
|
return jsonify({"error": "Authentication required"}), 401
|
|
|
|
payload = verify_token_via_broker(raw_token)
|
|
if not payload or not payload.get("valid"):
|
|
return jsonify({"error": "Invalid or expired token"}), 401
|
|
|
|
request.current_user = {
|
|
"user_id": payload.get("user_id"),
|
|
"token_id": payload.get("token_id"),
|
|
"scopes": payload.get("scopes", []),
|
|
}
|
|
return f(*args, **kwargs)
|
|
return decorated
|
|
|
|
|
|
def admin_required(f):
|
|
"""Decorator: require admin scope on the authenticated token."""
|
|
@wraps(f)
|
|
def decorated(*args, **kwargs):
|
|
auth_header = request.headers.get("Authorization", "")
|
|
api_key = request.headers.get("X-API-Key", "")
|
|
|
|
raw_token = ""
|
|
if auth_header.startswith("Bearer "):
|
|
raw_token = auth_header.split(" ", 1)[1]
|
|
elif api_key:
|
|
raw_token = api_key
|
|
|
|
if not raw_token:
|
|
return jsonify({"error": "Authentication required"}), 401
|
|
|
|
payload = verify_token_via_broker(raw_token)
|
|
if not payload or not payload.get("valid"):
|
|
return jsonify({"error": "Invalid or expired token"}), 401
|
|
|
|
scopes = payload.get("scopes", [])
|
|
if "admin" not in scopes and "ai_router_admin" not in scopes:
|
|
return jsonify({"error": "Admin access required"}), 403
|
|
|
|
request.current_user = {
|
|
"user_id": payload.get("user_id"),
|
|
"token_id": payload.get("token_id"),
|
|
"scopes": scopes,
|
|
}
|
|
return f(*args, **kwargs)
|
|
return decorated
|