Spaces:
Runtime error
Runtime error
File size: 2,751 Bytes
9e93b10 | 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 | import logging
import time
from typing import Any, Optional
from urllib.parse import quote
import jwt
from rexpro_ai.env import (
FORWARD_USER_INFO_HEADER_JWT,
FORWARD_USER_INFO_HEADER_JWT_EXPIRES_SECONDS,
FORWARD_USER_INFO_HEADER_JWT_SECRET,
FORWARD_USER_INFO_HEADER_USER_EMAIL,
FORWARD_USER_INFO_HEADER_USER_ID,
FORWARD_USER_INFO_HEADER_USER_NAME,
FORWARD_USER_INFO_HEADER_USER_ROLE,
)
log = logging.getLogger(__name__)
def _mint_forward_user_jwt(user: Any) -> str:
now = int(time.time())
payload = {
'sub': str(user.id),
'email': str(user.email),
'name': str(user.name),
'role': str(user.role),
'iss': 'rexpro-ai',
'iat': now,
'exp': now + FORWARD_USER_INFO_HEADER_JWT_EXPIRES_SECONDS,
}
return jwt.encode(payload, FORWARD_USER_INFO_HEADER_JWT_SECRET, algorithm='HS256')
def include_user_info_headers(headers: dict, user: Optional[Any] = None) -> dict:
"""
Forward user identity to external backends: signed JWT in
FORWARD_USER_INFO_HEADER_JWT if FORWARD_USER_INFO_HEADER_JWT_SECRET is set;
otherwise the legacy X-rexpro-ai-User-* headers.
"""
if user is None:
return headers
if FORWARD_USER_INFO_HEADER_JWT_SECRET:
try:
token = _mint_forward_user_jwt(user)
return {**headers, FORWARD_USER_INFO_HEADER_JWT: token}
except Exception:
log.exception(
'Failed to mint %s; falling back to plain user-info headers.',
FORWARD_USER_INFO_HEADER_JWT,
)
return {
**headers,
FORWARD_USER_INFO_HEADER_USER_NAME: quote(user.name, safe=' '),
FORWARD_USER_INFO_HEADER_USER_ID: user.id,
FORWARD_USER_INFO_HEADER_USER_EMAIL: user.email,
FORWARD_USER_INFO_HEADER_USER_ROLE: user.role,
}
def get_custom_headers(custom_headers: dict, user=None, metadata: dict = None) -> dict:
if not custom_headers or not isinstance(custom_headers, dict):
return {}
metadata = metadata or {}
template_vars = {
'{{CHAT_ID}}': metadata.get('chat_id', '') or '',
'{{MESSAGE_ID}}': metadata.get('message_id', '') or '',
'{{USER_ID}}': (user.id if user else '') or '',
'{{USER_NAME}}': (user.name if user else '') or '',
'{{USER_EMAIL}}': (user.email if user else '') or '',
'{{USER_ROLE}}': (user.role if user else '') or '',
}
parsed_headers = {}
for key, value in custom_headers.items():
if not isinstance(value, str):
value = str(value)
for token, val in template_vars.items():
value = value.replace(token, val)
parsed_headers[key] = value
return parsed_headers
|