mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-09 20:50:20 +08:00
feat: 优化批量余额查询和用户模型权限检查
- 添加批量余额查询并发限制,避免数据库连接池耗尽 - 支持余额加载 pending 状态和前端自动重试机制 - 添加用户可用模型 API,统一使用 AccessRestrictions - 修复用户表单编辑时数组引用共享导致的数据覆盖问题 - 添加数据库连接池配置说明到 .env.example
This commit is contained in:
@@ -62,3 +62,32 @@ ADMIN_PASSWORD=admin123456
|
||||
# 从发起请求到收到第一个字节的最大等待时间
|
||||
# 仅对流式请求生效,超时触发故障转移
|
||||
# STREAM_FIRST_BYTE_TIMEOUT=30.0
|
||||
|
||||
# ==================== 数据库连接池配置 ====================
|
||||
# 连接池大小直接影响并发能力,特别是流式请求场景
|
||||
# 每个流式请求会占用一个连接直到响应完成
|
||||
|
||||
# PostgreSQL 最大连接数(默认 100,需与 postgresql.conf 中 max_connections 匹配)
|
||||
# PG_MAX_CONNECTIONS=100
|
||||
|
||||
# 预留给管理工具的连接数(默认 10)
|
||||
# PG_RESERVED_CONNECTIONS=10
|
||||
|
||||
# 连接池大小(默认自动计算:(PG_MAX - RESERVED) / WORKERS / 2)
|
||||
# 高并发场景建议手动设置较大值
|
||||
# DB_POOL_SIZE=20
|
||||
|
||||
# 最大溢出连接数(默认等于 DB_POOL_SIZE)
|
||||
# 高峰期可临时创建的额外连接
|
||||
# DB_MAX_OVERFLOW=20
|
||||
|
||||
# 连接获取超时(默认 60 秒)
|
||||
# 超时后抛出 TimeoutError
|
||||
# DB_POOL_TIMEOUT=60
|
||||
|
||||
# 连接池使用率警告阈值(默认 70%)
|
||||
# DB_POOL_WARN_THRESHOLD=70
|
||||
|
||||
# 批量余额查询并发限制(默认自动计算:连接池的 40%)
|
||||
# 控制 Provider 余额查询的并发数,避免耗尽连接池
|
||||
# BATCH_BALANCE_CONCURRENCY=8
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import apiClient from './client'
|
||||
import type { ActivityHeatmap } from '@/types/activity'
|
||||
import type { TieredPricingConfig } from './endpoints/types'
|
||||
|
||||
export interface Profile {
|
||||
id: string // UUID
|
||||
@@ -235,6 +236,28 @@ export const meApi = {
|
||||
return response.data
|
||||
},
|
||||
|
||||
// 获取用户可用的模型列表
|
||||
async getAvailableModels(params?: {
|
||||
skip?: number
|
||||
limit?: number
|
||||
search?: string
|
||||
}): Promise<{
|
||||
models: Array<{
|
||||
id: string
|
||||
name: string
|
||||
display_name: string | null
|
||||
is_active: boolean
|
||||
default_price_per_request: number | null
|
||||
default_tiered_pricing: TieredPricingConfig | null
|
||||
supported_capabilities: string[] | null
|
||||
config: Record<string, any> | null
|
||||
}>
|
||||
total: number
|
||||
}> {
|
||||
const response = await apiClient.get('/api/users/me/available-models', { params })
|
||||
return response.data
|
||||
},
|
||||
|
||||
// 获取端点状态(不包含敏感信息)
|
||||
async getEndpointStatus(): Promise<any[]> {
|
||||
const response = await apiClient.get('/api/users/me/endpoint-status')
|
||||
|
||||
@@ -11,7 +11,7 @@ export interface PublicGlobalModel {
|
||||
display_name: string | null
|
||||
is_active: boolean
|
||||
// 阶梯计费配置
|
||||
default_tiered_pricing: TieredPricingConfig
|
||||
default_tiered_pricing: TieredPricingConfig | null
|
||||
default_price_per_request: number | null // 按次计费价格
|
||||
// Key 能力支持
|
||||
supported_capabilities: string[] | null
|
||||
|
||||
@@ -101,11 +101,15 @@ export function useFormDialog<E>(
|
||||
})
|
||||
|
||||
// 监听实体变化(编辑模式切换)
|
||||
watch(entity, (newEntity) => {
|
||||
if (newEntity && isOpen()) {
|
||||
// 注意:不使用 deep: true,只在实体引用变化时触发(如从 user A 切换到 user B)
|
||||
// 使用 deep: true 会导致实体的深层属性变化时也触发,覆盖用户正在编辑的数据
|
||||
watch(entity, (newEntity, oldEntity) => {
|
||||
// 只在实体引用真正变化时(如切换用户)才重新加载
|
||||
// 避免深层属性变化时意外触发
|
||||
if (newEntity && isOpen() && newEntity !== oldEntity) {
|
||||
loadData()
|
||||
}
|
||||
}, { immediate: true, deep: true })
|
||||
}, { immediate: true })
|
||||
|
||||
return {
|
||||
isEditMode,
|
||||
|
||||
@@ -457,6 +457,7 @@ function loadUserData() {
|
||||
if (!props.user) return
|
||||
formNonce.value = createFieldNonce()
|
||||
passwordFocused.value = false
|
||||
// 创建数组副本,避免与 props 数据共享引用
|
||||
form.value = {
|
||||
username: props.user.username,
|
||||
password: '',
|
||||
@@ -466,9 +467,9 @@ function loadUserData() {
|
||||
role: props.user.role,
|
||||
unlimited: props.user.quota_usd == null,
|
||||
is_active: props.user.is_active ?? true,
|
||||
allowed_providers: props.user.allowed_providers || [],
|
||||
allowed_api_formats: props.user.allowed_api_formats || [],
|
||||
allowed_models: props.user.allowed_models || []
|
||||
allowed_providers: [...(props.user.allowed_providers || [])],
|
||||
allowed_api_formats: [...(props.user.allowed_api_formats || [])],
|
||||
allowed_models: [...(props.user.allowed_models || [])]
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -143,9 +143,17 @@
|
||||
</div>
|
||||
</TableCell>
|
||||
<TableCell class="py-3.5">
|
||||
<!-- 余额正在加载中 -->
|
||||
<div
|
||||
v-if="provider.ops_configured && isBalanceLoading(provider.id)"
|
||||
class="flex items-center gap-1.5 text-xs text-muted-foreground"
|
||||
>
|
||||
<Loader2 class="h-3 w-3 animate-spin" />
|
||||
<span>加载中...</span>
|
||||
</div>
|
||||
<!-- 显示从上游 API 查询的余额 -->
|
||||
<div
|
||||
v-if="provider.ops_configured && getProviderBalance(provider.id)"
|
||||
v-else-if="provider.ops_configured && getProviderBalance(provider.id)"
|
||||
class="flex items-center gap-2 text-xs"
|
||||
>
|
||||
<!-- 余额文字 -->
|
||||
@@ -445,9 +453,17 @@
|
||||
>
|
||||
{{ formatBillingType(provider.billing_type || 'pay_as_you_go') }}
|
||||
</Badge>
|
||||
<!-- 余额加载中 -->
|
||||
<span
|
||||
v-if="provider.ops_configured && isBalanceLoading(provider.id)"
|
||||
class="text-muted-foreground flex items-center gap-1"
|
||||
>
|
||||
<Loader2 class="h-3 w-3 animate-spin" />
|
||||
加载中...
|
||||
</span>
|
||||
<!-- 余额(从上游 API 查询) -->
|
||||
<span
|
||||
v-if="provider.ops_configured && getProviderBalance(provider.id)"
|
||||
v-else-if="provider.ops_configured && getProviderBalance(provider.id)"
|
||||
class="text-muted-foreground"
|
||||
>
|
||||
余额 <span class="font-semibold text-foreground/90">{{ formatBalanceDisplay(getProviderBalance(provider.id)) }}</span>
|
||||
@@ -586,7 +602,8 @@ import {
|
||||
Trash2,
|
||||
ChevronDown,
|
||||
Power,
|
||||
KeyRound
|
||||
KeyRound,
|
||||
Loader2
|
||||
} from 'lucide-vue-next'
|
||||
import Button from '@/components/ui/button.vue'
|
||||
import Badge from '@/components/ui/badge.vue'
|
||||
@@ -737,17 +754,73 @@ async function loadBalances() {
|
||||
// 检查是否有新的请求已经开始,如果有则丢弃当前结果
|
||||
if (currentVersion !== balanceLoadVersion) return
|
||||
|
||||
// 将成功的结果存入缓存
|
||||
// 收集需要重试的 provider IDs
|
||||
const pendingProviderIds: string[] = []
|
||||
|
||||
// 将结果存入缓存(包括 pending 状态)
|
||||
for (const [providerId, result] of Object.entries(results)) {
|
||||
if (result.status === 'success') {
|
||||
// 存入缓存:success, auth_expired (带有效数据), pending
|
||||
if (result.status === 'success' || result.status === 'auth_expired' || result.status === 'pending') {
|
||||
balanceCache.value[providerId] = result
|
||||
}
|
||||
// 收集 pending 状态的 provider,稍后重试
|
||||
if (result.status === 'pending') {
|
||||
pendingProviderIds.push(providerId)
|
||||
}
|
||||
}
|
||||
|
||||
// 如果有 pending 状态的 provider,3秒后自动重试
|
||||
if (pendingProviderIds.length > 0) {
|
||||
const timerId = setTimeout(() => {
|
||||
pendingTimers.delete(timerId)
|
||||
// 检查版本号,确保没有新的加载请求
|
||||
if (currentVersion === balanceLoadVersion) {
|
||||
retryPendingBalances(pendingProviderIds, currentVersion, 0)
|
||||
}
|
||||
}, 3000)
|
||||
pendingTimers.add(timerId)
|
||||
}
|
||||
} catch (e) {
|
||||
console.warn('[loadBalances] 加载余额数据失败:', e)
|
||||
}
|
||||
}
|
||||
|
||||
// 重试加载 pending 状态的余额
|
||||
const MAX_BALANCE_RETRIES = 3
|
||||
|
||||
// 追踪待处理的定时器,用于组件卸载时清理
|
||||
const pendingTimers = new Set<ReturnType<typeof setTimeout>>()
|
||||
|
||||
async function retryPendingBalances(providerIds: string[], loadVersion: number, retryCount: number) {
|
||||
try {
|
||||
const results = await batchQueryBalance(providerIds)
|
||||
const stillPending: string[] = []
|
||||
|
||||
for (const [providerId, result] of Object.entries(results)) {
|
||||
if (result.status !== 'pending') {
|
||||
balanceCache.value[providerId] = result
|
||||
} else {
|
||||
stillPending.push(providerId)
|
||||
}
|
||||
}
|
||||
|
||||
// 如果还有 pending 且未达到最大重试次数,继续重试(指数退避)
|
||||
if (stillPending.length > 0 && retryCount < MAX_BALANCE_RETRIES) {
|
||||
const delay = 3000 * Math.pow(1.5, retryCount) // 3s, 4.5s, 6.75s
|
||||
const timerId = setTimeout(() => {
|
||||
pendingTimers.delete(timerId)
|
||||
// 检查版本号,确保没有新的加载请求
|
||||
if (loadVersion === balanceLoadVersion) {
|
||||
retryPendingBalances(stillPending, loadVersion, retryCount + 1)
|
||||
}
|
||||
}, delay)
|
||||
pendingTimers.add(timerId)
|
||||
}
|
||||
} catch (e) {
|
||||
console.warn('[retryPendingBalances] 重试加载余额失败:', e)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 类型守卫:检查是否为 BalanceInfo(简化版)
|
||||
* 只检查余额显示所需的字段,完整的 BalanceInfo 还包含 total_granted, total_used, expires_at, extra
|
||||
@@ -785,6 +858,10 @@ function getProviderBalanceError(providerId: string): { status: string; message:
|
||||
if (!result) {
|
||||
return null
|
||||
}
|
||||
// pending 状态不是错误,正在加载中
|
||||
if (result.status === 'pending') {
|
||||
return null
|
||||
}
|
||||
// 认证失败或过期
|
||||
if (result.status === 'auth_failed' || result.status === 'auth_expired') {
|
||||
return {
|
||||
@@ -802,6 +879,12 @@ function getProviderBalanceError(providerId: string): { status: string; message:
|
||||
return null
|
||||
}
|
||||
|
||||
// 检查余额是否正在加载中
|
||||
function isBalanceLoading(providerId: string): boolean {
|
||||
const result = balanceCache.value[providerId]
|
||||
return result?.status === 'pending'
|
||||
}
|
||||
|
||||
// 获取 provider 的签到信息(从 extra 字段)
|
||||
function getProviderCheckin(providerId: string): { success: boolean | null; message: string } | null {
|
||||
const result = balanceCache.value[providerId]
|
||||
@@ -1098,5 +1181,8 @@ onUnmounted(() => {
|
||||
if (tickInterval) {
|
||||
clearInterval(tickInterval)
|
||||
}
|
||||
// 清理余额重试的待处理定时器
|
||||
pendingTimers.forEach(clearTimeout)
|
||||
pendingTimers.clear()
|
||||
})
|
||||
</script>
|
||||
|
||||
@@ -920,6 +920,7 @@ function openCreateDialog() {
|
||||
}
|
||||
|
||||
function editUser(user: any) {
|
||||
// 创建数组副本,避免与 store 数据共享引用
|
||||
editingUser.value = {
|
||||
id: user.id,
|
||||
username: user.username,
|
||||
@@ -927,9 +928,9 @@ function editUser(user: any) {
|
||||
quota_usd: user.quota_usd,
|
||||
role: user.role,
|
||||
is_active: user.is_active,
|
||||
allowed_providers: user.allowed_providers || [],
|
||||
allowed_api_formats: user.allowed_api_formats || [],
|
||||
allowed_models: user.allowed_models || []
|
||||
allowed_providers: [...(user.allowed_providers || [])],
|
||||
allowed_api_formats: [...(user.allowed_api_formats || [])],
|
||||
allowed_models: [...(user.allowed_models || [])]
|
||||
}
|
||||
showUserFormDialog.value = true
|
||||
}
|
||||
|
||||
@@ -357,10 +357,7 @@ import {
|
||||
Pagination,
|
||||
RefreshButton,
|
||||
} from '@/components/ui'
|
||||
import {
|
||||
getPublicGlobalModels,
|
||||
type PublicGlobalModel,
|
||||
} from '@/api/public-models'
|
||||
import { type PublicGlobalModel } from '@/api/public-models'
|
||||
import { meApi } from '@/api/me'
|
||||
import {
|
||||
getUserConfigurableCapabilities,
|
||||
@@ -515,8 +512,9 @@ watch([searchQuery, capabilityFilters], () => {
|
||||
async function loadModels() {
|
||||
loading.value = true
|
||||
try {
|
||||
const response = await getPublicGlobalModels({ limit: 1000 })
|
||||
models.value = response.models || []
|
||||
// 使用用户认证端点,只获取用户有权限使用的模型
|
||||
const response = await meApi.getAvailableModels({ limit: 1000 })
|
||||
models.value = (response.models || []) as PublicGlobalModel[]
|
||||
} catch (err: any) {
|
||||
log.error('加载模型失败:', err)
|
||||
showError(err.response?.data?.detail || err.message, '加载模型失败')
|
||||
|
||||
+128
-2
@@ -19,11 +19,13 @@ from src.database import get_db
|
||||
from src.models.api import (
|
||||
ChangePasswordRequest,
|
||||
CreateMyApiKeyRequest,
|
||||
PublicGlobalModelListResponse,
|
||||
PublicGlobalModelResponse,
|
||||
UpdateApiKeyProvidersRequest,
|
||||
UpdatePreferencesRequest,
|
||||
UpdateProfileRequest,
|
||||
)
|
||||
from src.models.database import ApiKey, Provider, Usage, User
|
||||
from src.models.database import ApiKey, GlobalModel, Model, Provider, Usage, User
|
||||
from src.services.usage.service import UsageService
|
||||
from src.services.user.apikey import ApiKeyService
|
||||
from src.services.user.preference import PreferenceService
|
||||
@@ -261,6 +263,34 @@ async def list_available_providers(request: Request, db: Session = Depends(get_d
|
||||
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||
|
||||
|
||||
@router.get("/available-models")
|
||||
async def list_available_models(
|
||||
request: Request,
|
||||
skip: int = Query(0, ge=0, description="跳过记录数"),
|
||||
limit: int = Query(100, ge=1, le=1000, description="返回记录数限制"),
|
||||
search: Optional[str] = Query(None, description="搜索关键词"),
|
||||
db: Session = Depends(get_db),
|
||||
):
|
||||
"""
|
||||
获取用户可用的模型列表
|
||||
|
||||
根据用户权限返回可用的 GlobalModel 列表。
|
||||
- 管理员:可以看到所有活跃提供商的模型
|
||||
- 普通用户:只能看到关联提供商的模型
|
||||
|
||||
**查询参数**:
|
||||
- skip: 跳过的记录数,用于分页,默认 0
|
||||
- limit: 返回记录数限制,默认 100,范围 1-1000
|
||||
- search: 可选,搜索关键词,支持模糊匹配模型名称
|
||||
|
||||
**返回字段**:
|
||||
- models: 模型列表
|
||||
- total: 符合条件的模型总数
|
||||
"""
|
||||
adapter = ListAvailableModelsAdapter(skip=skip, limit=limit, search=search)
|
||||
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||
|
||||
|
||||
@router.get("/endpoint-status")
|
||||
async def get_endpoint_status(request: Request, db: Session = Depends(get_db)):
|
||||
"""
|
||||
@@ -981,13 +1011,109 @@ class GetMyActivityHeatmapAdapter(AuthenticatedApiAdapter):
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class ListAvailableModelsAdapter(AuthenticatedApiAdapter):
|
||||
"""获取用户可用模型列表的适配器"""
|
||||
|
||||
skip: int
|
||||
limit: int
|
||||
search: Optional[str]
|
||||
|
||||
async def handle(self, context): # type: ignore[override]
|
||||
from sqlalchemy import or_
|
||||
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
db = context.db
|
||||
user = context.user
|
||||
|
||||
# 使用 AccessRestrictions 类来处理限制(与 /v1/models 逻辑一致)
|
||||
restrictions = AccessRestrictions.from_api_key_and_user(api_key=None, user=user)
|
||||
|
||||
# 获取所有活跃的 Provider ID
|
||||
all_active_provider_ids = {
|
||||
p.id for p in db.query(Provider.id).filter(Provider.is_active == True).all()
|
||||
}
|
||||
|
||||
if not all_active_provider_ids:
|
||||
return {"models": [], "total": 0}
|
||||
|
||||
# 查询所有活跃的 GlobalModel 及其关联的 Model
|
||||
id_query = (
|
||||
db.query(GlobalModel.id, GlobalModel.name, Model.provider_id)
|
||||
.join(Model, Model.global_model_id == GlobalModel.id)
|
||||
.filter(
|
||||
and_(
|
||||
Model.provider_id.in_(all_active_provider_ids),
|
||||
Model.is_active == True,
|
||||
GlobalModel.is_active == True,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
# 搜索过滤
|
||||
if self.search:
|
||||
search_term = f"%{self.search}%"
|
||||
id_query = id_query.filter(
|
||||
or_(
|
||||
GlobalModel.name.ilike(search_term),
|
||||
GlobalModel.display_name.ilike(search_term),
|
||||
)
|
||||
)
|
||||
|
||||
# 获取所有匹配的记录
|
||||
all_matches = id_query.all()
|
||||
|
||||
# 应用访问限制过滤
|
||||
allowed_global_model_ids = set()
|
||||
for global_model_id, model_name, provider_id in all_matches:
|
||||
# 使用 AccessRestrictions.is_model_allowed 检查模型是否可访问
|
||||
# 它会同时检查 allowed_providers 和 allowed_models
|
||||
if restrictions.is_model_allowed(model_name, provider_id):
|
||||
allowed_global_model_ids.add(global_model_id)
|
||||
|
||||
# 统计总数
|
||||
total = len(allowed_global_model_ids)
|
||||
|
||||
if not allowed_global_model_ids:
|
||||
return {"models": [], "total": 0}
|
||||
|
||||
# 分页并获取完整的 GlobalModel 对象
|
||||
models = (
|
||||
db.query(GlobalModel)
|
||||
.filter(GlobalModel.id.in_(allowed_global_model_ids))
|
||||
.order_by(GlobalModel.name)
|
||||
.offset(self.skip)
|
||||
.limit(self.limit)
|
||||
.all()
|
||||
)
|
||||
|
||||
# 转换为响应格式(复用 PublicGlobalModelResponse schema)
|
||||
model_responses = [
|
||||
PublicGlobalModelResponse(
|
||||
id=gm.id,
|
||||
name=gm.name,
|
||||
display_name=gm.display_name,
|
||||
is_active=gm.is_active,
|
||||
default_price_per_request=gm.default_price_per_request,
|
||||
default_tiered_pricing=gm.default_tiered_pricing,
|
||||
supported_capabilities=gm.supported_capabilities,
|
||||
config=gm.config,
|
||||
)
|
||||
for gm in models
|
||||
]
|
||||
|
||||
logger.debug(f"用户 {user.email} 可用模型: {len(model_responses)} 个")
|
||||
return PublicGlobalModelListResponse(models=model_responses, total=total)
|
||||
|
||||
|
||||
class ListAvailableProvidersAdapter(AuthenticatedApiAdapter):
|
||||
"""获取可用提供商列表的适配器"""
|
||||
|
||||
async def handle(self, context): # type: ignore[override]
|
||||
from sqlalchemy.orm import selectinload
|
||||
|
||||
from src.models.database import Model, ProviderEndpoint
|
||||
from src.models.database import ProviderEndpoint
|
||||
|
||||
db = context.db
|
||||
|
||||
|
||||
@@ -5,12 +5,14 @@ Provider 操作服务
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from dataclasses import asdict
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.config import config
|
||||
from src.core.cache_service import CacheService
|
||||
from src.core.crypto import CryptoService
|
||||
from src.core.logger import logger
|
||||
@@ -33,6 +35,36 @@ from src.services.provider_ops.types import (
|
||||
BALANCE_CACHE_TTL = 86400
|
||||
|
||||
|
||||
def _get_batch_balance_concurrency() -> int:
|
||||
"""
|
||||
动态计算批量余额查询的并发限制
|
||||
|
||||
计算逻辑:
|
||||
1. 优先使用环境变量 BATCH_BALANCE_CONCURRENCY
|
||||
2. 否则根据连接池大小自动计算(取 40% 的连接池容量)
|
||||
3. 限制在 [3, 15] 范围内
|
||||
|
||||
Returns:
|
||||
并发限制数
|
||||
"""
|
||||
# 优先使用环境变量
|
||||
env_value = os.getenv("BATCH_BALANCE_CONCURRENCY")
|
||||
if env_value:
|
||||
try:
|
||||
return max(1, int(env_value))
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# 根据连接池大小自动计算
|
||||
# 连接池容量 = pool_size + max_overflow
|
||||
pool_capacity = config.db_pool_size + config.db_max_overflow
|
||||
|
||||
# 取 40% 的连接池容量,保留 60% 给其他请求
|
||||
# 最小 3(保证基本并发),最大 15(避免过度并发)
|
||||
calculated = int(pool_capacity * 0.4)
|
||||
return max(3, min(calculated, 15))
|
||||
|
||||
|
||||
class ProviderOpsService:
|
||||
"""
|
||||
Provider 操作服务
|
||||
@@ -366,6 +398,7 @@ class ProviderOpsService:
|
||||
self,
|
||||
provider_id: str,
|
||||
trigger_refresh: bool = True,
|
||||
allow_sync_query: bool = True,
|
||||
) -> ActionResult:
|
||||
"""
|
||||
查询余额(优先返回缓存,可触发异步刷新)
|
||||
@@ -373,6 +406,7 @@ class ProviderOpsService:
|
||||
Args:
|
||||
provider_id: Provider ID
|
||||
trigger_refresh: 是否触发后台异步刷新
|
||||
allow_sync_query: 缓存未命中时是否允许同步查询(False 时仅返回缓存或触发异步刷新)
|
||||
|
||||
Returns:
|
||||
操作结果(可能是缓存的)
|
||||
@@ -387,19 +421,43 @@ class ProviderOpsService:
|
||||
asyncio.create_task(self._refresh_balance_async(provider_id))
|
||||
return cached
|
||||
|
||||
# 没有缓存,同步查询一次(首次访问)
|
||||
logger.info(f"余额缓存未命中,同步查询: provider_id={provider_id}")
|
||||
return await self.query_balance(provider_id)
|
||||
# 没有缓存
|
||||
if allow_sync_query:
|
||||
# 同步查询一次(首次访问)
|
||||
logger.info(f"余额缓存未命中,同步查询: provider_id={provider_id}")
|
||||
return await self.query_balance(provider_id)
|
||||
else:
|
||||
# 仅触发异步刷新,立即返回
|
||||
logger.debug(f"余额缓存未命中,触发异步刷新: provider_id={provider_id}")
|
||||
asyncio.create_task(self._refresh_balance_async(provider_id))
|
||||
return ActionResult(
|
||||
status=ActionStatus.PENDING,
|
||||
action_type=ProviderActionType.QUERY_BALANCE,
|
||||
message="余额数据加载中,请稍后刷新",
|
||||
)
|
||||
|
||||
async def _refresh_balance_async(self, provider_id: str) -> None:
|
||||
"""后台异步刷新余额(使用独立的数据库 session)"""
|
||||
"""
|
||||
后台异步刷新余额(使用独立的数据库 session)
|
||||
|
||||
注意:这是一个后台任务,使用独立的短生命周期 session,
|
||||
避免长时间占用连接池资源。
|
||||
"""
|
||||
db = None
|
||||
try:
|
||||
# 后台任务需要创建独立的 session,因为原请求的 session 可能已关闭
|
||||
with create_session() as db:
|
||||
service = ProviderOpsService(db)
|
||||
await service.query_balance(provider_id)
|
||||
db = create_session()
|
||||
service = ProviderOpsService(db)
|
||||
await service.query_balance(provider_id)
|
||||
except Exception as e:
|
||||
logger.warning(f"异步刷新余额失败: provider_id={provider_id}, error={e}")
|
||||
finally:
|
||||
# 确保 session 被关闭,归还连接到连接池
|
||||
if db is not None:
|
||||
try:
|
||||
db.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _clear_balance_cache(self, provider_id: str) -> None:
|
||||
"""清除余额缓存(认证失败时调用)"""
|
||||
@@ -643,6 +701,8 @@ class ProviderOpsService:
|
||||
"""
|
||||
批量查询余额(优先返回缓存,后台异步刷新)
|
||||
|
||||
使用信号量限制并发数,避免数据库连接池耗尽。
|
||||
|
||||
Args:
|
||||
provider_ids: Provider ID 列表,None 表示查询所有已配置的
|
||||
|
||||
@@ -658,26 +718,36 @@ class ProviderOpsService:
|
||||
if p.config and p.config.get("provider_ops")
|
||||
]
|
||||
|
||||
# 并行查询,使用缓存优先策略
|
||||
tasks = [
|
||||
self.query_balance_with_cache(provider_id, trigger_refresh=True)
|
||||
for provider_id in provider_ids
|
||||
]
|
||||
results_list = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
if not provider_ids:
|
||||
return {}
|
||||
|
||||
results = {}
|
||||
for provider_id, result in zip(provider_ids, results_list):
|
||||
if isinstance(result, Exception):
|
||||
logger.warning(f"查询余额失败: provider_id={provider_id}, error={result}")
|
||||
results[provider_id] = ActionResult(
|
||||
status=ActionStatus.UNKNOWN_ERROR,
|
||||
action_type=ProviderActionType.QUERY_BALANCE,
|
||||
message=str(result),
|
||||
)
|
||||
else:
|
||||
results[provider_id] = result
|
||||
# 使用信号量限制并发数,避免同时发起过多请求耗尽连接池
|
||||
concurrency = _get_batch_balance_concurrency()
|
||||
semaphore = asyncio.Semaphore(concurrency)
|
||||
logger.debug(f"批量余额查询: providers={len(provider_ids)}, concurrency={concurrency}")
|
||||
|
||||
return results
|
||||
async def _query_with_limit(provider_id: str) -> tuple[str, ActionResult]:
|
||||
async with semaphore:
|
||||
try:
|
||||
# 批量查询时禁用同步查询,避免阻塞请求
|
||||
# 缓存未命中时会触发异步刷新,前端可稍后重试
|
||||
result = await self.query_balance_with_cache(
|
||||
provider_id, trigger_refresh=True, allow_sync_query=False
|
||||
)
|
||||
return provider_id, result
|
||||
except Exception as e:
|
||||
logger.warning(f"查询余额失败: provider_id={provider_id}, error={e}")
|
||||
return provider_id, ActionResult(
|
||||
status=ActionStatus.UNKNOWN_ERROR,
|
||||
action_type=ProviderActionType.QUERY_BALANCE,
|
||||
message=str(e),
|
||||
)
|
||||
|
||||
# 并行查询,但受信号量限制
|
||||
tasks = [_query_with_limit(provider_id) for provider_id in provider_ids]
|
||||
results_list = await asyncio.gather(*tasks)
|
||||
|
||||
return dict(results_list)
|
||||
|
||||
# ==================== 认证验证 ====================
|
||||
|
||||
|
||||
@@ -34,6 +34,7 @@ class ActionStatus(str, Enum):
|
||||
"""操作执行状态"""
|
||||
|
||||
SUCCESS = "success" # 成功
|
||||
PENDING = "pending" # 处理中(异步任务已触发,尚未完成)
|
||||
AUTH_FAILED = "auth_failed" # 认证失败
|
||||
AUTH_EXPIRED = "auth_expired" # 认证过期
|
||||
RATE_LIMITED = "rate_limited" # 频率限制
|
||||
|
||||
@@ -408,28 +408,28 @@ class UserService:
|
||||
"""获取用户可用的模型
|
||||
|
||||
通过 GlobalModel + Model 关联查询用户可用模型
|
||||
逻辑:用户可用提供商 -> Provider 的 Model 实现 -> 关联的 GlobalModel
|
||||
逻辑:使用 AccessRestrictions 统一处理 allowed_providers 和 allowed_models 限制
|
||||
"""
|
||||
# 获取用户可用的提供商
|
||||
if user.role == UserRole.ADMIN:
|
||||
# 管理员可以使用所有活动提供商
|
||||
provider_ids = [
|
||||
p.id for p in db.query(Provider.id).filter(Provider.is_active == True).all()
|
||||
]
|
||||
else:
|
||||
# 普通用户使用关联的提供商
|
||||
provider_ids = [p.id for p in user.providers]
|
||||
from src.api.base.models_service import AccessRestrictions
|
||||
|
||||
if not provider_ids:
|
||||
# 使用 AccessRestrictions 类来处理限制(与 /v1/models 逻辑一致)
|
||||
restrictions = AccessRestrictions.from_api_key_and_user(api_key=None, user=user)
|
||||
|
||||
# 获取所有活跃的 Provider ID
|
||||
all_active_provider_ids = [
|
||||
p.id for p in db.query(Provider.id).filter(Provider.is_active == True).all()
|
||||
]
|
||||
|
||||
if not all_active_provider_ids:
|
||||
return []
|
||||
|
||||
# 查询这些提供商的所有活跃 Model(关联 GlobalModel)
|
||||
models = (
|
||||
# 查询所有活跃的 Model(关联 GlobalModel)
|
||||
all_models = (
|
||||
db.query(Model)
|
||||
.join(GlobalModel, Model.global_model_id == GlobalModel.id)
|
||||
.filter(
|
||||
and_(
|
||||
Model.provider_id.in_(provider_ids),
|
||||
Model.provider_id.in_(all_active_provider_ids),
|
||||
Model.is_active == True,
|
||||
GlobalModel.is_active == True,
|
||||
)
|
||||
@@ -437,6 +437,14 @@ class UserService:
|
||||
.all()
|
||||
)
|
||||
|
||||
logger.debug(f"用户 {user.email} 可用模型: {len(models)} 个 (提供商数: {len(provider_ids)})")
|
||||
# 应用访问限制过滤
|
||||
filtered_models = []
|
||||
for model in all_models:
|
||||
model_name = model.global_model.name if model.global_model else model.provider_model_name
|
||||
# 使用 AccessRestrictions.is_model_allowed 检查模型是否可访问
|
||||
if restrictions.is_model_allowed(model_name, model.provider_id):
|
||||
filtered_models.append(model)
|
||||
|
||||
return models
|
||||
logger.debug(f"用户 {user.email} 可用模型: {len(filtered_models)} 个")
|
||||
|
||||
return filtered_models
|
||||
|
||||
Reference in New Issue
Block a user