forked from xiaohei/taiji-AI-PAD
771 lines
26 KiB
Python
771 lines
26 KiB
Python
"""
|
||
认证与权限管理路由
|
||
"""
|
||
|
||
from datetime import timedelta
|
||
from fastapi import APIRouter, Depends, HTTPException, status, Query
|
||
from sqlalchemy import select, and_
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from sqlalchemy.exc import IntegrityError
|
||
import secrets
|
||
import hashlib
|
||
from typing import Optional
|
||
|
||
from database import get_db
|
||
from models import (
|
||
User, Channel, APIKey, Balance, TenantCustomAgentQuota,
|
||
PlatformAgentQuota, TenantModelKey, ResourceAllocation, ModelProvider
|
||
)
|
||
from app.auth import (
|
||
authenticate_user,
|
||
create_access_token,
|
||
create_refresh_token,
|
||
get_password_hash,
|
||
verify_password,
|
||
require_auth,
|
||
)
|
||
from app.schemas import (
|
||
LoginRequest,
|
||
SuccessResponse,
|
||
TokenResponse,
|
||
PasswordChangeRequest,
|
||
APIKeyInfo,
|
||
RegenerateAPIKeyResponse,
|
||
UserCreate,
|
||
)
|
||
from app.email_verification import verify_code, send_and_store_verification_code, check_rate_limit
|
||
from app.agent_manager_client import get_agent_manager_client, AgentManagerError
|
||
from app.litellm_client import get_litellm_client, LiteLLMClientError
|
||
from config import settings
|
||
import uuid
|
||
from datetime import datetime
|
||
from sqlalchemy import and_
|
||
import structlog
|
||
|
||
logger = structlog.get_logger(__name__)
|
||
|
||
router = APIRouter(prefix="/api/auth", tags=["认证"])
|
||
|
||
|
||
def _mask_api_key(key: str) -> str:
|
||
"""隐藏API密钥的中间部分"""
|
||
if len(key) <= 12:
|
||
return key[:4] + "..." + key[-4:]
|
||
return key[:8] + "..." + key[-4:]
|
||
|
||
|
||
@router.post("/login", response_model=SuccessResponse)
|
||
async def login(req: LoginRequest, db: AsyncSession = Depends(get_db)):
|
||
"""
|
||
用户/渠道/管理员/供应商登录
|
||
|
||
支持的角色：
|
||
- user: 租户用户
|
||
- channel: 渠道管理员
|
||
- billing_admin: 计费管理员
|
||
- operations_admin: 运营管理员
|
||
- admin: 管理员
|
||
- super_admin: 超级管理员
|
||
- provider: 供应商管理员
|
||
"""
|
||
# 根据角色查找用户
|
||
if req.role == "channel":
|
||
# 渠道管理员登录
|
||
result = await db.execute(select(Channel).where(Channel.email == req.email))
|
||
entity = result.scalar_one_or_none()
|
||
|
||
if not entity or not verify_password(req.password, entity.password_hash):
|
||
raise HTTPException(
|
||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||
detail="邮箱或密码错误"
|
||
)
|
||
|
||
# 更新最后登录时间
|
||
from datetime import datetime
|
||
entity.last_login_at = datetime.utcnow()
|
||
await db.commit()
|
||
|
||
# 创建JWT token
|
||
token_data = {
|
||
"sub": str(entity.id),
|
||
"email": entity.email,
|
||
"role": "channel_admin",
|
||
"channelId": str(entity.id),
|
||
}
|
||
access_token = create_access_token(data=token_data)
|
||
refresh_token = create_refresh_token(data=token_data)
|
||
|
||
return SuccessResponse(
|
||
data={
|
||
"token": access_token,
|
||
"refreshToken": refresh_token,
|
||
"user": {
|
||
"id": str(entity.id),
|
||
"name": entity.name,
|
||
"email": entity.email,
|
||
"role": "channel_admin",
|
||
"channelId": str(entity.id),
|
||
}
|
||
}
|
||
)
|
||
|
||
else:
|
||
# 用户/管理员/供应商登录
|
||
result = await db.execute(select(User).where(User.email == req.email))
|
||
user = result.scalar_one_or_none()
|
||
|
||
if not user:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||
detail="邮箱或密码错误"
|
||
)
|
||
|
||
# 验证密码（兼容两种密码字段）
|
||
password_hash = user.password_hash or user.hashed_password
|
||
if not password_hash or not verify_password(req.password, password_hash):
|
||
raise HTTPException(
|
||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||
detail="邮箱或密码错误"
|
||
)
|
||
|
||
# 验证用户状态
|
||
if hasattr(user, 'status') and user.status == "inactive":
|
||
raise HTTPException(
|
||
status_code=status.HTTP_403_FORBIDDEN,
|
||
detail="账户已被禁用"
|
||
)
|
||
|
||
# 验证角色
|
||
user_role = user.role
|
||
|
||
# 角色验证逻辑
|
||
valid_roles = {
|
||
"super_admin": ["super_admin"],
|
||
"admin": ["admin", "super_admin"],
|
||
"billing_admin": ["billing_admin", "admin", "super_admin"],
|
||
"operations_admin": ["operations_admin", "admin", "super_admin"],
|
||
"user": ["user"],
|
||
"provider": ["provider_admin"],
|
||
}
|
||
|
||
allowed_roles = valid_roles.get(req.role, [])
|
||
if user_role not in allowed_roles:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_403_FORBIDDEN,
|
||
detail=f"权限不足，当前角色: {user_role}"
|
||
)
|
||
|
||
# 更新最后登录时间
|
||
from datetime import datetime
|
||
user.last_login_at = datetime.utcnow()
|
||
await db.commit()
|
||
|
||
# 创建JWT token
|
||
token_data = {
|
||
"sub": str(user.id),
|
||
"email": user.email,
|
||
"role": user_role,
|
||
"channelId": str(user.channel_id) if user.channel_id else None,
|
||
}
|
||
access_token = create_access_token(data=token_data)
|
||
refresh_token = create_refresh_token(data=token_data)
|
||
|
||
return SuccessResponse(
|
||
data={
|
||
"token": access_token,
|
||
"refreshToken": refresh_token,
|
||
"user": {
|
||
"id": str(user.id),
|
||
"name": user.name or user.full_name,
|
||
"email": user.email,
|
||
"role": user_role,
|
||
"channelId": str(user.channel_id) if user.channel_id else None,
|
||
}
|
||
}
|
||
)
|
||
|
||
|
||
@router.post("/logout", response_model=SuccessResponse)
|
||
async def logout(
|
||
principal: dict = Depends(require_auth),
|
||
db: AsyncSession = Depends(get_db)
|
||
):
|
||
"""
|
||
用户登出
|
||
|
||
将当前Token加入黑名单
|
||
"""
|
||
from datetime import datetime, timedelta
|
||
from app.token_blacklist import add_token_to_blacklist
|
||
import uuid as uuid_module
|
||
|
||
user_id = principal.get("user_id")
|
||
claims = principal.get("claims", {})
|
||
|
||
# 确保 user_id 是有效的 UUID 字符串
|
||
if user_id:
|
||
try:
|
||
# 验证并转换为标准 UUID 字符串格式
|
||
user_id = str(uuid_module.UUID(str(user_id)))
|
||
except (ValueError, TypeError):
|
||
# 如果无法转换，使用 claims 中的 sub
|
||
sub = claims.get("sub")
|
||
if sub:
|
||
try:
|
||
user_id = str(uuid_module.UUID(str(sub)))
|
||
except (ValueError, TypeError):
|
||
user_id = None
|
||
|
||
if not user_id:
|
||
# 如果无法获取有效的 user_id，直接返回成功（前端会清除本地 token）
|
||
return SuccessResponse(message="登出成功")
|
||
|
||
# 获取Token的过期时间（从claims中提取）
|
||
exp = claims.get("exp")
|
||
if exp:
|
||
expires_at = datetime.utcfromtimestamp(exp)
|
||
else:
|
||
# 默认7天后过期
|
||
expires_at = datetime.utcnow() + timedelta(days=7)
|
||
|
||
# 生成token的唯一标识（使用sub + 时间戳）
|
||
token_jti = f"logout_{user_id}_{datetime.utcnow().timestamp()}"
|
||
|
||
try:
|
||
# 添加到黑名单
|
||
await add_token_to_blacklist(
|
||
token_jti=token_jti,
|
||
user_id=user_id,
|
||
expires_at=expires_at,
|
||
reason="logout",
|
||
db=db,
|
||
)
|
||
except Exception as e:
|
||
# 即使黑名单添加失败，也返回成功（前端会清除本地 token）
|
||
import structlog
|
||
logger = structlog.get_logger(__name__)
|
||
logger.warning("failed_to_add_token_to_blacklist", error=str(e), user_id=user_id)
|
||
|
||
return SuccessResponse(message="登出成功")
|
||
|
||
|
||
@router.post("/refresh", response_model=SuccessResponse)
|
||
async def refresh_token_endpoint(
|
||
principal: dict = Depends(require_auth),
|
||
db: AsyncSession = Depends(get_db)
|
||
):
|
||
"""
|
||
刷新访问令牌
|
||
|
||
使用 Refresh Token 获取新的 Access Token
|
||
"""
|
||
user_id = principal.get("user_id")
|
||
claims = principal.get("claims", {})
|
||
|
||
# 验证是否为 refresh token（可选，如果前端确保传入的是 refresh token）
|
||
token_type = claims.get("type")
|
||
# 为了向后兼容，不强制要求 type 为 refresh
|
||
|
||
# 查询用户信息
|
||
result = await db.execute(select(User).where(User.id == user_id))
|
||
user = result.scalar_one_or_none()
|
||
|
||
if not user:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||
detail="用户不存在"
|
||
)
|
||
|
||
# 生成新的 access token 和 refresh token
|
||
token_data = {
|
||
"sub": str(user.id),
|
||
"email": user.email,
|
||
"role": user.role,
|
||
"channelId": str(user.channel_id) if user.channel_id else None,
|
||
}
|
||
new_access_token = create_access_token(data=token_data)
|
||
new_refresh_token = create_refresh_token(data=token_data)
|
||
|
||
return SuccessResponse(
|
||
data={
|
||
"token": new_access_token,
|
||
"refreshToken": new_refresh_token,
|
||
}
|
||
)
|
||
|
||
|
||
@router.put("/password", response_model=SuccessResponse)
|
||
async def change_password(
|
||
req: PasswordChangeRequest,
|
||
principal: dict = Depends(require_auth),
|
||
db: AsyncSession = Depends(get_db)
|
||
):
|
||
"""
|
||
修改密码
|
||
"""
|
||
user_id = principal.get("user_id")
|
||
|
||
# 查询用户
|
||
result = await db.execute(select(User).where(User.id == user_id))
|
||
user = result.scalar_one_or_none()
|
||
|
||
if not user:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_404_NOT_FOUND,
|
||
detail="用户不存在"
|
||
)
|
||
|
||
# 验证旧密码
|
||
password_hash = user.password_hash or user.hashed_password
|
||
if not verify_password(req.old_password, password_hash):
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail="旧密码错误"
|
||
)
|
||
|
||
# 更新密码
|
||
new_hash = get_password_hash(req.new_password)
|
||
user.password_hash = new_hash
|
||
user.hashed_password = new_hash # 兼容旧字段
|
||
|
||
await db.commit()
|
||
|
||
return SuccessResponse(message="密码修改成功")
|
||
|
||
|
||
@router.get("/keys/info", response_model=SuccessResponse)
|
||
async def get_api_key_info(
|
||
principal: dict = Depends(require_auth),
|
||
db: AsyncSession = Depends(get_db)
|
||
):
|
||
"""
|
||
获取服务终结点和API密钥信息
|
||
"""
|
||
user_id = principal.get("user_id")
|
||
|
||
# 查询用户的API密钥
|
||
result = await db.execute(
|
||
select(APIKey)
|
||
.where(APIKey.user_id == user_id)
|
||
.where(APIKey.is_active == True)
|
||
.order_by(APIKey.created_at.desc())
|
||
)
|
||
api_key = result.scalar_one_or_none()
|
||
|
||
if not api_key:
|
||
# 创建新的API密钥
|
||
raw_key = f"sk-{secrets.token_urlsafe(32)}"
|
||
key_hash = get_password_hash(raw_key)
|
||
|
||
new_key = APIKey(
|
||
user_id=user_id,
|
||
api_key_hash=key_hash,
|
||
api_key_prefix=raw_key[:8],
|
||
name="默认密钥",
|
||
key_hash=key_hash,
|
||
prefix=raw_key[:8],
|
||
is_active=True,
|
||
)
|
||
db.add(new_key)
|
||
await db.commit()
|
||
await db.refresh(new_key)
|
||
|
||
api_key = new_key
|
||
masked_key = _mask_api_key(raw_key)
|
||
else:
|
||
# 隐藏现有密钥
|
||
masked_key = api_key.api_key_prefix + "..." + "xxxx"
|
||
|
||
endpoint = f"https://api.taiji-ai.com/v1" # 可以从配置读取
|
||
|
||
return SuccessResponse(
|
||
data={
|
||
"endpoint": endpoint,
|
||
"apiKey": masked_key,
|
||
"createdAt": api_key.created_at.isoformat(),
|
||
"lastUsed": api_key.last_used.isoformat() if api_key.last_used else None,
|
||
}
|
||
)
|
||
|
||
|
||
@router.post("/keys/regenerate", response_model=SuccessResponse)
|
||
async def regenerate_api_key(
|
||
principal: dict = Depends(require_auth),
|
||
db: AsyncSession = Depends(get_db)
|
||
):
|
||
"""
|
||
重新生成API密钥
|
||
"""
|
||
user_id = principal.get("user_id")
|
||
|
||
# 禁用所有旧密钥
|
||
result = await db.execute(
|
||
select(APIKey).where(APIKey.user_id == user_id)
|
||
)
|
||
old_keys = result.scalars().all()
|
||
|
||
for key in old_keys:
|
||
key.is_active = False
|
||
|
||
# 生成新密钥
|
||
raw_key = f"sk-{secrets.token_urlsafe(32)}"
|
||
key_hash = get_password_hash(raw_key)
|
||
|
||
new_key = APIKey(
|
||
user_id=user_id,
|
||
api_key_hash=key_hash,
|
||
api_key_prefix=raw_key[:8],
|
||
name="重新生成的密钥",
|
||
key_hash=key_hash,
|
||
prefix=raw_key[:8],
|
||
is_active=True,
|
||
)
|
||
db.add(new_key)
|
||
await db.commit()
|
||
|
||
return SuccessResponse(
|
||
data={
|
||
"apiKey": raw_key, # 完整显示新密钥
|
||
"message": "旧密钥已失效",
|
||
}
|
||
)
|
||
|
||
|
||
@router.post("/register/send-code", response_model=SuccessResponse)
|
||
async def send_verification_code_endpoint(
|
||
email: str = Query(..., description="邮箱地址"),
|
||
db: AsyncSession = Depends(get_db)
|
||
):
|
||
"""
|
||
发送邮箱验证码
|
||
|
||
在用户注册前，先调用此接口发送验证码到邮箱
|
||
|
||
频率限制：同一邮箱60秒内只能发送一次
|
||
"""
|
||
# 1. 检查发送频率限制
|
||
can_send, remaining_seconds = await check_rate_limit(email)
|
||
if not can_send:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
||
detail=f"请等待{remaining_seconds}秒后再重新发送验证码"
|
||
)
|
||
|
||
# 2. 检查邮箱是否已存在
|
||
result = await db.execute(select(User).where(User.email == email))
|
||
existing_user = result.scalar_one_or_none()
|
||
|
||
if existing_user:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail="该邮箱已被注册"
|
||
)
|
||
|
||
# 3. 发送验证码
|
||
code = await send_and_store_verification_code(email)
|
||
if not code:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||
detail="验证码发送失败，请稍后重试"
|
||
)
|
||
|
||
return SuccessResponse(
|
||
message="验证码已发送到您的邮箱，请查收"
|
||
)
|
||
|
||
|
||
@router.post("/register", response_model=SuccessResponse)
|
||
async def register(req: UserCreate, db: AsyncSession = Depends(get_db)):
|
||
"""
|
||
用户注册接口
|
||
|
||
允许用户自由注册，创建普通用户账户
|
||
注册成功后自动分配默认资源：
|
||
- 分配到 taiji 渠道 (channelId: b415e70b-8d37-481c-b229-bc3b7871607b)
|
||
- 自定义 agent 配额：2 CPU, 2 GB 内存
|
||
- 平台 agent 各1个
|
||
- 供应商所有模型
|
||
- 余额：20元
|
||
"""
|
||
# 1. 先检查邮箱是否已存在（不消耗验证码）
|
||
result = await db.execute(select(User).where(User.email == req.email))
|
||
existing_user = result.scalar_one_or_none()
|
||
|
||
if existing_user:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail="该邮箱已被注册"
|
||
)
|
||
|
||
# 2. 检查用户名是否已存在（username 是必填字段）
|
||
result = await db.execute(select(User).where(User.username == req.username))
|
||
existing_username = result.scalar_one_or_none()
|
||
if existing_username:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail="该用户名已被使用"
|
||
)
|
||
|
||
# 3. 最后验证邮箱验证码（验证成功后会消耗验证码）
|
||
is_valid = await verify_code(req.email, req.verification_code)
|
||
if not is_valid:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail="验证码错误或已过期"
|
||
)
|
||
|
||
# taiji 渠道 ID
|
||
TAIJI_CHANNEL_ID = uuid.UUID("b415e70b-8d37-481c-b229-bc3b7871607b")
|
||
|
||
# 验证渠道存在
|
||
channel_result = await db.execute(
|
||
select(Channel).where(Channel.id == TAIJI_CHANNEL_ID)
|
||
)
|
||
channel = channel_result.scalar_one_or_none()
|
||
if not channel:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||
detail="渠道不存在，请联系管理员"
|
||
)
|
||
|
||
# 创建新用户
|
||
password_hash = get_password_hash(req.password)
|
||
username = req.username or req.email.split("@")[0]
|
||
name = req.full_name or username
|
||
|
||
new_user = User(
|
||
name=name,
|
||
email=req.email,
|
||
password_hash=password_hash,
|
||
hashed_password=password_hash, # 兼容字段
|
||
username=username,
|
||
full_name=req.full_name or name,
|
||
role="user", # 默认角色为普通用户
|
||
channel_id=TAIJI_CHANNEL_ID, # 分配到 taiji 渠道
|
||
status="active",
|
||
is_active=True,
|
||
is_admin=False,
|
||
balance=0,
|
||
credit_limit=0,
|
||
eu_balance=0,
|
||
total_eu_consumed=0,
|
||
)
|
||
|
||
try:
|
||
db.add(new_user)
|
||
await db.flush() # 获取 user.id，这里会触发唯一约束检查
|
||
except IntegrityError as e:
|
||
await db.rollback()
|
||
error_str = str(e.orig) if e.orig else str(e)
|
||
if "email" in error_str.lower() or "users_email_key" in error_str.lower():
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail="该邮箱已被注册"
|
||
)
|
||
elif "username" in error_str.lower() or "users_username_key" in error_str.lower():
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail="该用户名已被使用"
|
||
)
|
||
else:
|
||
logger.error(f"注册时数据库约束冲突: {e}")
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail="注册信息与现有用户冲突，请更换邮箱或用户名"
|
||
)
|
||
|
||
user_id = new_user.id
|
||
|
||
try:
|
||
# 1. 创建余额记录，初始余额 20 元
|
||
balance = Balance(
|
||
user_id=user_id,
|
||
eu_balance=20.0
|
||
)
|
||
db.add(balance)
|
||
|
||
# 2. 分配自定义 Agent 配额：2 CPU, 2 GB 内存
|
||
custom_agent_quota = TenantCustomAgentQuota(
|
||
tenant_id=user_id,
|
||
cpu_quota=2.0,
|
||
memory_quota=2.0,
|
||
cpu_used=0.0,
|
||
memory_used=0.0,
|
||
agent_count=0
|
||
)
|
||
db.add(custom_agent_quota)
|
||
|
||
# 3. 获取所有平台 Agent 模板，为每个模板分配 1 个配额
|
||
try:
|
||
agent_manager_client = get_agent_manager_client()
|
||
platform_templates = await agent_manager_client.list_platform_templates()
|
||
|
||
for template in platform_templates:
|
||
template_name = template.template
|
||
|
||
# 获取渠道的配额记录（需要同步更新 pod_used）
|
||
channel_quota_result = await db.execute(
|
||
select(PlatformAgentQuota)
|
||
.where(
|
||
and_(
|
||
PlatformAgentQuota.target_id == TAIJI_CHANNEL_ID,
|
||
PlatformAgentQuota.target_type == "channel",
|
||
PlatformAgentQuota.template_name == template_name
|
||
)
|
||
)
|
||
.with_for_update() # 行锁
|
||
)
|
||
channel_quota = channel_quota_result.scalar_one_or_none()
|
||
|
||
# 如果渠道没有该模板的配额，跳过
|
||
if not channel_quota:
|
||
logger.warning(f"渠道没有平台 Agent '{template_name}' 的配额，跳过分配")
|
||
continue
|
||
|
||
# 检查渠道剩余配额是否足够
|
||
remaining = channel_quota.pod_quota - (channel_quota.pod_used or 0)
|
||
if remaining < 1:
|
||
logger.warning(f"渠道平台 Agent '{template_name}' 配额不足，跳过分配")
|
||
continue
|
||
|
||
# 创建平台 Agent 配额记录
|
||
platform_quota = PlatformAgentQuota(
|
||
target_id=user_id,
|
||
target_type="tenant",
|
||
template_name=template_name,
|
||
pod_quota=1, # 每个平台 agent 分配 1 个配额
|
||
pod_used=0,
|
||
allocated_by=None, # 系统自动分配
|
||
allocated_at=datetime.utcnow()
|
||
)
|
||
db.add(platform_quota)
|
||
|
||
# Bug 修复：更新渠道的 pod_used（分配给租户的配额视为渠道已使用的配额）
|
||
channel_quota.pod_used = (channel_quota.pod_used or 0) + 1
|
||
|
||
# 同时创建 ResourceAllocation 记录（兼容旧逻辑）
|
||
allocation = ResourceAllocation(
|
||
target_id=user_id,
|
||
target_type="tenant",
|
||
resource_type="agent",
|
||
resource_id=template_name,
|
||
quantity=1
|
||
)
|
||
db.add(allocation)
|
||
except (AgentManagerError, Exception) as e:
|
||
logger.warning(f"获取平台 Agent 模板失败，跳过平台 Agent 分配: {e}")
|
||
|
||
# 4. 分配所有供应商模型
|
||
providers_result = await db.execute(
|
||
select(ModelProvider).where(ModelProvider.is_active == True)
|
||
)
|
||
providers = providers_result.scalars().all()
|
||
|
||
if channel.litellm_team_id:
|
||
litellm_client = get_litellm_client()
|
||
|
||
for provider in providers:
|
||
# 为每个供应商的每个模型创建 TenantModelKey
|
||
for model_name in provider.supported_models:
|
||
try:
|
||
# 在 LiteLLM 中创建 Key
|
||
key = await litellm_client.generate_key(
|
||
team_id=channel.litellm_team_id,
|
||
models=[model_name],
|
||
rpm_limit=provider.rpm or 60,
|
||
tpm_limit=provider.tpm or 10000,
|
||
max_budget=500.0, # 默认预算
|
||
budget_duration="monthly",
|
||
key_name=f"tenant-{user_id}-{model_name}",
|
||
metadata={
|
||
"tenant_id": str(user_id),
|
||
"tenant_name": name,
|
||
"channel_id": str(TAIJI_CHANNEL_ID),
|
||
"channel_name": channel.name,
|
||
"model": model_name,
|
||
}
|
||
)
|
||
|
||
# 加密存储 Key
|
||
encrypted_key = litellm_client.encrypt_key(key.key)
|
||
|
||
# 保存到数据库
|
||
tenant_key = TenantModelKey(
|
||
tenant_id=user_id,
|
||
channel_id=TAIJI_CHANNEL_ID,
|
||
model_name=model_name,
|
||
litellm_key_id=key.key,
|
||
litellm_key_hash=encrypted_key,
|
||
rpm_limit=provider.rpm or 60,
|
||
tpm_limit=provider.tpm or 10000,
|
||
max_budget=500.0,
|
||
budget_duration="monthly",
|
||
status="active",
|
||
)
|
||
db.add(tenant_key)
|
||
|
||
# 同时记录到 ResourceAllocation
|
||
model_allocation = ResourceAllocation(
|
||
target_id=user_id,
|
||
target_type="tenant",
|
||
resource_type="model",
|
||
resource_id=str(provider.id),
|
||
rpm=provider.rpm or 60,
|
||
tpm=provider.tpm or 10000,
|
||
)
|
||
db.add(model_allocation)
|
||
except (LiteLLMClientError, Exception) as e:
|
||
logger.warning(f"为模型 {model_name} 创建 LiteLLM Key 失败: {e}")
|
||
# 继续处理其他模型
|
||
else:
|
||
logger.warning(f"渠道 {TAIJI_CHANNEL_ID} 未配置 LiteLLM team，跳过模型分配")
|
||
|
||
# 提交所有更改
|
||
await db.commit()
|
||
await db.refresh(new_user)
|
||
|
||
logger.info(
|
||
f"用户注册成功并分配默认资源",
|
||
user_id=str(user_id),
|
||
email=req.email,
|
||
channel_id=str(TAIJI_CHANNEL_ID)
|
||
)
|
||
|
||
except Exception as e:
|
||
await db.rollback()
|
||
logger.error(f"用户注册失败: {e}")
|
||
raise HTTPException(
|
||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||
detail=f"注册失败: {str(e)}"
|
||
)
|
||
|
||
# 创建JWT token，自动登录
|
||
token_data = {
|
||
"sub": str(new_user.id),
|
||
"email": new_user.email,
|
||
"role": new_user.role,
|
||
"user_id": str(new_user.id),
|
||
"channelId": str(TAIJI_CHANNEL_ID),
|
||
}
|
||
|
||
# 创建 Access Token（短期有效）和 Refresh Token（长期有效）
|
||
access_token = create_access_token(data=token_data)
|
||
refresh_token = create_refresh_token(data=token_data)
|
||
|
||
return SuccessResponse(
|
||
data={
|
||
"token": access_token,
|
||
"refreshToken": refresh_token,
|
||
"user": {
|
||
"id": str(new_user.id),
|
||
"name": new_user.name,
|
||
"email": new_user.email,
|
||
"username": new_user.username,
|
||
"role": new_user.role,
|
||
"channelId": str(TAIJI_CHANNEL_ID),
|
||
}
|
||
},
|
||
message="注册成功"
|
||
)
|
||
|