forked from xiaohei/taiji-AI-PAD
509 lines
14 KiB
Python
509 lines
14 KiB
Python
"""
|
||
渠道合作伙伴API路由
|
||
"""
|
||
|
||
from datetime import datetime, timedelta
|
||
from typing import List, Optional
|
||
from fastapi import APIRouter, Depends, HTTPException, Query, status
|
||
from sqlalchemy import select, func, and_, desc, or_
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
import uuid
|
||
|
||
from database import get_db
|
||
from models import (
|
||
User, Channel, Agent, ResourceAllocation,
|
||
BillingRecord, RechargeRecord, Application, ModelProvider
|
||
)
|
||
from app.auth import require_auth, get_password_hash
|
||
from app.schemas import (
|
||
SuccessResponse,
|
||
TenantCreateRequest,
|
||
AllocateResourcesRequest,
|
||
UpdateBillingRequest,
|
||
RechargeTenantRequest,
|
||
RechargeTenantResponse,
|
||
SetCreditLimitRequest,
|
||
SetCreditLimitResponse,
|
||
ResourceApplicationRequest,
|
||
ChannelBillingResponse,
|
||
)
|
||
|
||
router = APIRouter(prefix="/api/channel", tags=["渠道合作伙伴"])
|
||
|
||
|
||
def _verify_channel_permission(principal: dict):
|
||
"""验证渠道管理员权限"""
|
||
role = principal.get("claims", {}).get("role")
|
||
if role != "channel_admin":
|
||
raise HTTPException(
|
||
status_code=status.HTTP_403_FORBIDDEN,
|
||
detail="需要渠道管理员权限"
|
||
)
|
||
|
||
|
||
# ============= 租户管理 =============
|
||
|
||
@router.get("/tenants", response_model=SuccessResponse)
|
||
async def list_tenants(
|
||
principal: dict = Depends(require_auth),
|
||
db: AsyncSession = Depends(get_db)
|
||
):
|
||
"""
|
||
获取渠道下的租户列表
|
||
"""
|
||
_verify_channel_permission(principal)
|
||
channel_id = principal.get("claims", {}).get("channelId")
|
||
|
||
if not channel_id:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail="无法获取渠道ID"
|
||
)
|
||
|
||
result = await db.execute(
|
||
select(User).where(User.channel_id == channel_id)
|
||
)
|
||
tenants = result.scalars().all()
|
||
|
||
data = [
|
||
{
|
||
"id": str(tenant.id),
|
||
"name": tenant.name,
|
||
"email": tenant.email,
|
||
"subscriptionTier": tenant.subscription_tier,
|
||
"balance": float(tenant.balance),
|
||
"creditLimit": float(tenant.credit_limit),
|
||
"status": tenant.status,
|
||
"createdAt": tenant.created_at.isoformat(),
|
||
}
|
||
for tenant in tenants
|
||
]
|
||
|
||
return SuccessResponse(data={"tenants": data})
|
||
|
||
|
||
@router.post("/tenants/create", response_model=SuccessResponse)
|
||
async def create_tenant(
|
||
req: TenantCreateRequest,
|
||
principal: dict = Depends(require_auth),
|
||
db: AsyncSession = Depends(get_db)
|
||
):
|
||
"""
|
||
创建租户
|
||
"""
|
||
_verify_channel_permission(principal)
|
||
channel_id = principal.get("claims", {}).get("channelId")
|
||
|
||
if not channel_id:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail="无法获取渠道ID"
|
||
)
|
||
|
||
# 检查邮箱是否已存在
|
||
result = await db.execute(
|
||
select(User).where(User.email == req.email)
|
||
)
|
||
if result.scalar_one_or_none():
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail="邮箱已被使用"
|
||
)
|
||
|
||
# 创建租户
|
||
password_hash = get_password_hash(req.password)
|
||
tenant = User(
|
||
name=req.name,
|
||
email=req.email,
|
||
password_hash=password_hash,
|
||
hashed_password=password_hash, # 兼容
|
||
username=req.email.split("@")[0],
|
||
full_name=req.name,
|
||
role="user",
|
||
channel_id=channel_id,
|
||
subscription_tier=req.subscriptionTier,
|
||
status="active",
|
||
balance=0,
|
||
credit_limit=0,
|
||
)
|
||
|
||
db.add(tenant)
|
||
await db.commit()
|
||
await db.refresh(tenant)
|
||
|
||
return SuccessResponse(
|
||
data={
|
||
"id": str(tenant.id),
|
||
"name": tenant.name,
|
||
"email": tenant.email,
|
||
},
|
||
message="租户创建成功"
|
||
)
|
||
|
||
|
||
@router.put("/tenants/{tenant_id}/resources", response_model=SuccessResponse)
|
||
async def allocate_tenant_resources(
|
||
tenant_id: str,
|
||
req: AllocateResourcesRequest,
|
||
principal: dict = Depends(require_auth),
|
||
db: AsyncSession = Depends(get_db)
|
||
):
|
||
"""
|
||
为租户分配资源
|
||
"""
|
||
_verify_channel_permission(principal)
|
||
channel_id = principal.get("claims", {}).get("channelId")
|
||
|
||
# 验证租户属于该渠道
|
||
result = await db.execute(
|
||
select(User).where(
|
||
and_(
|
||
User.id == tenant_id,
|
||
User.channel_id == channel_id
|
||
)
|
||
)
|
||
)
|
||
tenant = result.scalar_one_or_none()
|
||
|
||
if not tenant:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_404_NOT_FOUND,
|
||
detail="租户不存在或不属于该渠道"
|
||
)
|
||
|
||
# 删除现有资源分配
|
||
await db.execute(
|
||
select(ResourceAllocation).where(
|
||
and_(
|
||
ResourceAllocation.target_id == tenant_id,
|
||
ResourceAllocation.target_type == "tenant"
|
||
)
|
||
)
|
||
)
|
||
|
||
# 分配Agent资源
|
||
for agent_alloc in req.agents:
|
||
allocation = ResourceAllocation(
|
||
target_id=tenant_id,
|
||
target_type="tenant",
|
||
resource_type="agent",
|
||
resource_id=agent_alloc.agentId,
|
||
quantity=agent_alloc.quantity,
|
||
)
|
||
db.add(allocation)
|
||
|
||
# 分配模型资源
|
||
for model_alloc in req.models:
|
||
# 查找模型供应商
|
||
result = await db.execute(
|
||
select(ModelProvider).where(ModelProvider.name == model_alloc.modelName)
|
||
)
|
||
model_provider = result.scalar_one_or_none()
|
||
|
||
if model_provider:
|
||
allocation = ResourceAllocation(
|
||
target_id=tenant_id,
|
||
target_type="tenant",
|
||
resource_type="model",
|
||
resource_id=str(model_provider.id),
|
||
rpm=model_alloc.rpm,
|
||
tpm=model_alloc.tpm,
|
||
)
|
||
db.add(allocation)
|
||
|
||
await db.commit()
|
||
|
||
return SuccessResponse(message="资源分配成功")
|
||
|
||
|
||
@router.put("/tenants/{tenant_id}/billing", response_model=SuccessResponse)
|
||
async def update_tenant_billing(
|
||
tenant_id: str,
|
||
req: UpdateBillingRequest,
|
||
principal: dict = Depends(require_auth),
|
||
db: AsyncSession = Depends(get_db)
|
||
):
|
||
"""
|
||
更新租户计费设置
|
||
"""
|
||
_verify_channel_permission(principal)
|
||
channel_id = principal.get("claims", {}).get("channelId")
|
||
|
||
# 验证租户
|
||
result = await db.execute(
|
||
select(User).where(
|
||
and_(
|
||
User.id == tenant_id,
|
||
User.channel_id == channel_id
|
||
)
|
||
)
|
||
)
|
||
tenant = result.scalar_one_or_none()
|
||
|
||
if not tenant:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_404_NOT_FOUND,
|
||
detail="租户不存在"
|
||
)
|
||
|
||
# 更新计费设置
|
||
tenant.subscription_tier = req.subscriptionTier
|
||
tenant.discount = req.discount
|
||
|
||
await db.commit()
|
||
|
||
return SuccessResponse(message="计费设置更新成功")
|
||
|
||
|
||
@router.post("/tenants/{tenant_id}/recharge", response_model=SuccessResponse)
|
||
async def recharge_tenant(
|
||
tenant_id: str,
|
||
req: RechargeTenantRequest,
|
||
principal: dict = Depends(require_auth),
|
||
db: AsyncSession = Depends(get_db)
|
||
):
|
||
"""
|
||
为租户充值
|
||
"""
|
||
_verify_channel_permission(principal)
|
||
channel_id = principal.get("claims", {}).get("channelId")
|
||
|
||
# 验证租户
|
||
result = await db.execute(
|
||
select(User).where(
|
||
and_(
|
||
User.id == tenant_id,
|
||
User.channel_id == channel_id
|
||
)
|
||
)
|
||
)
|
||
tenant = result.scalar_one_or_none()
|
||
|
||
if not tenant:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_404_NOT_FOUND,
|
||
detail="租户不存在"
|
||
)
|
||
|
||
# 更新余额
|
||
tenant.balance = float(tenant.balance) + req.amount
|
||
|
||
# 创建充值记录
|
||
recharge = RechargeRecord(
|
||
user_id=tenant_id,
|
||
channel_id=channel_id,
|
||
amount=req.amount,
|
||
payment_method="channel_recharge",
|
||
status="success",
|
||
order_id=f"CH{datetime.utcnow().strftime('%Y%m%d%H%M%S')}{uuid.uuid4().hex[:8]}",
|
||
completed_at=datetime.utcnow(),
|
||
)
|
||
|
||
db.add(recharge)
|
||
await db.commit()
|
||
await db.refresh(tenant)
|
||
|
||
return SuccessResponse(
|
||
data={
|
||
"tenantId": str(tenant.id),
|
||
"newBalance": float(tenant.balance),
|
||
"rechargeAmount": req.amount,
|
||
}
|
||
)
|
||
|
||
|
||
@router.put("/tenants/{tenant_id}/credit", response_model=SuccessResponse)
|
||
async def set_tenant_credit_limit(
|
||
tenant_id: str,
|
||
req: SetCreditLimitRequest,
|
||
principal: dict = Depends(require_auth),
|
||
db: AsyncSession = Depends(get_db)
|
||
):
|
||
"""
|
||
设置租户授信额度
|
||
"""
|
||
_verify_channel_permission(principal)
|
||
channel_id = principal.get("claims", {}).get("channelId")
|
||
|
||
# 验证租户
|
||
result = await db.execute(
|
||
select(User).where(
|
||
and_(
|
||
User.id == tenant_id,
|
||
User.channel_id == channel_id
|
||
)
|
||
)
|
||
)
|
||
tenant = result.scalar_one_or_none()
|
||
|
||
if not tenant:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_404_NOT_FOUND,
|
||
detail="租户不存在"
|
||
)
|
||
|
||
# 更新授信额度
|
||
tenant.credit_limit = req.creditLimit
|
||
|
||
await db.commit()
|
||
|
||
return SuccessResponse(
|
||
data={
|
||
"tenantId": str(tenant.id),
|
||
"creditLimit": float(req.creditLimit),
|
||
}
|
||
)
|
||
|
||
|
||
# ============= 资源申请 =============
|
||
|
||
@router.post("/resources/apply", response_model=SuccessResponse)
|
||
async def apply_for_resources(
|
||
req: ResourceApplicationRequest,
|
||
principal: dict = Depends(require_auth),
|
||
db: AsyncSession = Depends(get_db)
|
||
):
|
||
"""
|
||
申请资源（模型或Agent）
|
||
"""
|
||
_verify_channel_permission(principal)
|
||
channel_id = principal.get("claims", {}).get("channelId")
|
||
|
||
if not channel_id:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail="无法获取渠道ID"
|
||
)
|
||
|
||
# 创建申请
|
||
application = Application(
|
||
channel_id=channel_id,
|
||
type=req.type,
|
||
model_name=req.modelName,
|
||
rpm=req.rpm,
|
||
tpm=req.tpm,
|
||
agent_type=req.agentType,
|
||
quantity=req.quantity,
|
||
reason=req.reason,
|
||
status="pending",
|
||
)
|
||
|
||
db.add(application)
|
||
await db.commit()
|
||
await db.refresh(application)
|
||
|
||
return SuccessResponse(
|
||
data={"id": str(application.id), "status": "pending"},
|
||
message="申请已提交，等待审批"
|
||
)
|
||
|
||
|
||
# ============= 计费统计 =============
|
||
|
||
@router.get("/billing/stats", response_model=SuccessResponse)
|
||
async def get_channel_billing_stats(
|
||
startTime: str = Query(...),
|
||
endTime: str = Query(...),
|
||
tenantName: Optional[str] = Query(None),
|
||
minCalls: Optional[int] = Query(None),
|
||
maxCalls: Optional[int] = Query(None),
|
||
export: Optional[str] = Query(None, pattern="^(excel|csv|pdf)$"),
|
||
principal: dict = Depends(require_auth),
|
||
db: AsyncSession = Depends(get_db)
|
||
):
|
||
"""
|
||
获取渠道计费统计
|
||
"""
|
||
_verify_channel_permission(principal)
|
||
channel_id = principal.get("claims", {}).get("channelId")
|
||
|
||
# 解析时间
|
||
start_dt = datetime.fromisoformat(startTime.replace("Z", "+00:00"))
|
||
end_dt = datetime.fromisoformat(endTime.replace("Z", "+00:00"))
|
||
|
||
# 查询渠道下的租户
|
||
tenants_result = await db.execute(
|
||
select(User).where(User.channel_id == channel_id)
|
||
)
|
||
tenants = {str(t.id): t for t in tenants_result.scalars().all()}
|
||
tenant_ids = list(tenants.keys())
|
||
|
||
# 租户统计
|
||
tenant_stats_result = await db.execute(
|
||
select(
|
||
BillingRecord.tenant_id,
|
||
func.count(BillingRecord.id).label("calls"),
|
||
func.sum(BillingRecord.eu).label("total_eu"),
|
||
func.sum(BillingRecord.cost).label("total_cost"),
|
||
)
|
||
.where(
|
||
and_(
|
||
BillingRecord.tenant_id.in_(tenant_ids),
|
||
BillingRecord.timestamp >= start_dt,
|
||
BillingRecord.timestamp <= end_dt,
|
||
)
|
||
)
|
||
.group_by(BillingRecord.tenant_id)
|
||
)
|
||
|
||
tenant_stats = []
|
||
for row in tenant_stats_result.all():
|
||
tenant = tenants.get(str(row.tenant_id))
|
||
if tenant:
|
||
tenant_stats.append({
|
||
"tenantId": str(row.tenant_id),
|
||
"tenantName": tenant.name,
|
||
"calls": row.calls,
|
||
"totalEU": float(row.total_eu or 0),
|
||
"totalCost": float(row.total_cost or 0),
|
||
})
|
||
|
||
# 调用记录
|
||
records_result = await db.execute(
|
||
select(BillingRecord)
|
||
.where(
|
||
and_(
|
||
BillingRecord.tenant_id.in_(tenant_ids),
|
||
BillingRecord.timestamp >= start_dt,
|
||
BillingRecord.timestamp <= end_dt,
|
||
)
|
||
)
|
||
.order_by(desc(BillingRecord.timestamp))
|
||
.limit(100)
|
||
)
|
||
|
||
call_records = []
|
||
for record in records_result.scalars().all():
|
||
tenant = tenants.get(str(record.tenant_id))
|
||
if tenant:
|
||
call_records.append({
|
||
"id": str(record.id),
|
||
"timestamp": record.timestamp.isoformat(),
|
||
"tenantName": tenant.name,
|
||
"agentName": record.agent_name,
|
||
"duration": record.duration,
|
||
"eu": record.eu,
|
||
"cost": float(record.cost),
|
||
})
|
||
|
||
# 如果是导出请求
|
||
if export:
|
||
file_url = f"https://exports.taiji-ai.com/{channel_id}/{export}/billing_{datetime.utcnow().strftime('%Y%m%d%H%M%S')}.{export}"
|
||
expires_at = (datetime.utcnow() + timedelta(hours=24)).isoformat()
|
||
|
||
return SuccessResponse(
|
||
data={
|
||
"fileUrl": file_url,
|
||
"format": export,
|
||
"expiresAt": expires_at,
|
||
}
|
||
)
|
||
|
||
return SuccessResponse(
|
||
data={
|
||
"tenantStats": tenant_stats,
|
||
"callRecords": call_records,
|
||
}
|
||
)
|
||
|