feat: 优化批量余额查询和用户模型权限检查

- 添加批量余额查询并发限制,避免数据库连接池耗尽
- 支持余额加载 pending 状态和前端自动重试机制
- 添加用户可用模型 API,统一使用 AccessRestrictions
- 修复用户表单编辑时数组引用共享导致的数据覆盖问题
- 添加数据库连接池配置说明到 .env.example
This commit is contained in:
fawney19
2026-01-28 01:01:05 +08:00
parent 8e0695e9d9
commit 1e0255d0ed
12 changed files with 411 additions and 64 deletions

View File

@@ -62,3 +62,32 @@ ADMIN_PASSWORD=admin123456
# 从发起请求到收到第一个字节的最大等待时间 # 从发起请求到收到第一个字节的最大等待时间
# 仅对流式请求生效,超时触发故障转移 # 仅对流式请求生效,超时触发故障转移
# STREAM_FIRST_BYTE_TIMEOUT=30.0 # 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

View File

@@ -1,5 +1,6 @@
import apiClient from './client' import apiClient from './client'
import type { ActivityHeatmap } from '@/types/activity' import type { ActivityHeatmap } from '@/types/activity'
import type { TieredPricingConfig } from './endpoints/types'
export interface Profile { export interface Profile {
id: string // UUID id: string // UUID
@@ -235,6 +236,28 @@ export const meApi = {
return response.data 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[]> { async getEndpointStatus(): Promise<any[]> {
const response = await apiClient.get('/api/users/me/endpoint-status') const response = await apiClient.get('/api/users/me/endpoint-status')

View File

@@ -11,7 +11,7 @@ export interface PublicGlobalModel {
display_name: string | null display_name: string | null
is_active: boolean is_active: boolean
// 阶梯计费配置 // 阶梯计费配置
default_tiered_pricing: TieredPricingConfig default_tiered_pricing: TieredPricingConfig | null
default_price_per_request: number | null // 按次计费价格 default_price_per_request: number | null // 按次计费价格
// Key 能力支持 // Key 能力支持
supported_capabilities: string[] | null supported_capabilities: string[] | null

View File

@@ -101,11 +101,15 @@ export function useFormDialog<E>(
}) })
// 监听实体变化(编辑模式切换) // 监听实体变化(编辑模式切换)
watch(entity, (newEntity) => { // 注意:不使用 deep: true只在实体引用变化时触发如从 user A 切换到 user B
if (newEntity && isOpen()) { // 使用 deep: true 会导致实体的深层属性变化时也触发,覆盖用户正在编辑的数据
watch(entity, (newEntity, oldEntity) => {
// 只在实体引用真正变化时(如切换用户)才重新加载
// 避免深层属性变化时意外触发
if (newEntity && isOpen() && newEntity !== oldEntity) {
loadData() loadData()
} }
}, { immediate: true, deep: true }) }, { immediate: true })
return { return {
isEditMode, isEditMode,

View File

@@ -457,6 +457,7 @@ function loadUserData() {
if (!props.user) return if (!props.user) return
formNonce.value = createFieldNonce() formNonce.value = createFieldNonce()
passwordFocused.value = false passwordFocused.value = false
// 创建数组副本,避免与 props 数据共享引用
form.value = { form.value = {
username: props.user.username, username: props.user.username,
password: '', password: '',
@@ -466,9 +467,9 @@ function loadUserData() {
role: props.user.role, role: props.user.role,
unlimited: props.user.quota_usd == null, unlimited: props.user.quota_usd == null,
is_active: props.user.is_active ?? true, is_active: props.user.is_active ?? true,
allowed_providers: props.user.allowed_providers || [], allowed_providers: [...(props.user.allowed_providers || [])],
allowed_api_formats: props.user.allowed_api_formats || [], allowed_api_formats: [...(props.user.allowed_api_formats || [])],
allowed_models: props.user.allowed_models || [] allowed_models: [...(props.user.allowed_models || [])]
} }
} }

View File

@@ -143,9 +143,17 @@
</div> </div>
</TableCell> </TableCell>
<TableCell class="py-3.5"> <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 查询的余额 --> <!-- 显示从上游 API 查询的余额 -->
<div <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" class="flex items-center gap-2 text-xs"
> >
<!-- 余额文字 --> <!-- 余额文字 -->
@@ -445,9 +453,17 @@
> >
{{ formatBillingType(provider.billing_type || 'pay_as_you_go') }} {{ formatBillingType(provider.billing_type || 'pay_as_you_go') }}
</Badge> </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 查询 --> <!-- 余额从上游 API 查询 -->
<span <span
v-if="provider.ops_configured && getProviderBalance(provider.id)" v-else-if="provider.ops_configured && getProviderBalance(provider.id)"
class="text-muted-foreground" class="text-muted-foreground"
> >
余额 <span class="font-semibold text-foreground/90">{{ formatBalanceDisplay(getProviderBalance(provider.id)) }}</span> 余额 <span class="font-semibold text-foreground/90">{{ formatBalanceDisplay(getProviderBalance(provider.id)) }}</span>
@@ -586,7 +602,8 @@ import {
Trash2, Trash2,
ChevronDown, ChevronDown,
Power, Power,
KeyRound KeyRound,
Loader2
} from 'lucide-vue-next' } from 'lucide-vue-next'
import Button from '@/components/ui/button.vue' import Button from '@/components/ui/button.vue'
import Badge from '@/components/ui/badge.vue' import Badge from '@/components/ui/badge.vue'
@@ -737,17 +754,73 @@ async function loadBalances() {
// 检查是否有新的请求已经开始,如果有则丢弃当前结果 // 检查是否有新的请求已经开始,如果有则丢弃当前结果
if (currentVersion !== balanceLoadVersion) return if (currentVersion !== balanceLoadVersion) return
// 将成功的结果存入缓存 // 收集需要重试的 provider IDs
const pendingProviderIds: string[] = []
// 将结果存入缓存(包括 pending 状态)
for (const [providerId, result] of Object.entries(results)) { 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 balanceCache.value[providerId] = result
} }
// 收集 pending 状态的 provider稍后重试
if (result.status === 'pending') {
pendingProviderIds.push(providerId)
}
}
// 如果有 pending 状态的 provider3秒后自动重试
if (pendingProviderIds.length > 0) {
const timerId = setTimeout(() => {
pendingTimers.delete(timerId)
// 检查版本号,确保没有新的加载请求
if (currentVersion === balanceLoadVersion) {
retryPendingBalances(pendingProviderIds, currentVersion, 0)
}
}, 3000)
pendingTimers.add(timerId)
} }
} catch (e) { } catch (e) {
console.warn('[loadBalances] 加载余额数据失败:', 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简化版
* 只检查余额显示所需的字段,完整的 BalanceInfo 还包含 total_granted, total_used, expires_at, extra * 只检查余额显示所需的字段,完整的 BalanceInfo 还包含 total_granted, total_used, expires_at, extra
@@ -785,6 +858,10 @@ function getProviderBalanceError(providerId: string): { status: string; message:
if (!result) { if (!result) {
return null return null
} }
// pending 状态不是错误,正在加载中
if (result.status === 'pending') {
return null
}
// 认证失败或过期 // 认证失败或过期
if (result.status === 'auth_failed' || result.status === 'auth_expired') { if (result.status === 'auth_failed' || result.status === 'auth_expired') {
return { return {
@@ -802,6 +879,12 @@ function getProviderBalanceError(providerId: string): { status: string; message:
return null return null
} }
// 检查余额是否正在加载中
function isBalanceLoading(providerId: string): boolean {
const result = balanceCache.value[providerId]
return result?.status === 'pending'
}
// 获取 provider 的签到信息(从 extra 字段) // 获取 provider 的签到信息(从 extra 字段)
function getProviderCheckin(providerId: string): { success: boolean | null; message: string } | null { function getProviderCheckin(providerId: string): { success: boolean | null; message: string } | null {
const result = balanceCache.value[providerId] const result = balanceCache.value[providerId]
@@ -1098,5 +1181,8 @@ onUnmounted(() => {
if (tickInterval) { if (tickInterval) {
clearInterval(tickInterval) clearInterval(tickInterval)
} }
// 清理余额重试的待处理定时器
pendingTimers.forEach(clearTimeout)
pendingTimers.clear()
}) })
</script> </script>

View File

@@ -920,6 +920,7 @@ function openCreateDialog() {
} }
function editUser(user: any) { function editUser(user: any) {
// 创建数组副本,避免与 store 数据共享引用
editingUser.value = { editingUser.value = {
id: user.id, id: user.id,
username: user.username, username: user.username,
@@ -927,9 +928,9 @@ function editUser(user: any) {
quota_usd: user.quota_usd, quota_usd: user.quota_usd,
role: user.role, role: user.role,
is_active: user.is_active, is_active: user.is_active,
allowed_providers: user.allowed_providers || [], allowed_providers: [...(user.allowed_providers || [])],
allowed_api_formats: user.allowed_api_formats || [], allowed_api_formats: [...(user.allowed_api_formats || [])],
allowed_models: user.allowed_models || [] allowed_models: [...(user.allowed_models || [])]
} }
showUserFormDialog.value = true showUserFormDialog.value = true
} }

View File

@@ -357,10 +357,7 @@ import {
Pagination, Pagination,
RefreshButton, RefreshButton,
} from '@/components/ui' } from '@/components/ui'
import { import { type PublicGlobalModel } from '@/api/public-models'
getPublicGlobalModels,
type PublicGlobalModel,
} from '@/api/public-models'
import { meApi } from '@/api/me' import { meApi } from '@/api/me'
import { import {
getUserConfigurableCapabilities, getUserConfigurableCapabilities,
@@ -515,8 +512,9 @@ watch([searchQuery, capabilityFilters], () => {
async function loadModels() { async function loadModels() {
loading.value = true loading.value = true
try { 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) { } catch (err: any) {
log.error('加载模型失败:', err) log.error('加载模型失败:', err)
showError(err.response?.data?.detail || err.message, '加载模型失败') showError(err.response?.data?.detail || err.message, '加载模型失败')

View File

@@ -19,11 +19,13 @@ from src.database import get_db
from src.models.api import ( from src.models.api import (
ChangePasswordRequest, ChangePasswordRequest,
CreateMyApiKeyRequest, CreateMyApiKeyRequest,
PublicGlobalModelListResponse,
PublicGlobalModelResponse,
UpdateApiKeyProvidersRequest, UpdateApiKeyProvidersRequest,
UpdatePreferencesRequest, UpdatePreferencesRequest,
UpdateProfileRequest, 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.usage.service import UsageService
from src.services.user.apikey import ApiKeyService from src.services.user.apikey import ApiKeyService
from src.services.user.preference import PreferenceService 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) 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") @router.get("/endpoint-status")
async def get_endpoint_status(request: Request, db: Session = Depends(get_db)): async def get_endpoint_status(request: Request, db: Session = Depends(get_db)):
""" """
@@ -981,13 +1011,109 @@ class GetMyActivityHeatmapAdapter(AuthenticatedApiAdapter):
return result 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): class ListAvailableProvidersAdapter(AuthenticatedApiAdapter):
"""获取可用提供商列表的适配器""" """获取可用提供商列表的适配器"""
async def handle(self, context): # type: ignore[override] async def handle(self, context): # type: ignore[override]
from sqlalchemy.orm import selectinload from sqlalchemy.orm import selectinload
from src.models.database import Model, ProviderEndpoint from src.models.database import ProviderEndpoint
db = context.db db = context.db

View File

@@ -5,12 +5,14 @@ Provider 操作服务
""" """
import asyncio import asyncio
import os
from dataclasses import asdict from dataclasses import asdict
from datetime import datetime, timezone from datetime import datetime, timezone
from typing import Any, Dict, List, Optional from typing import Any, Dict, List, Optional
from sqlalchemy.orm import Session from sqlalchemy.orm import Session
from src.config import config
from src.core.cache_service import CacheService from src.core.cache_service import CacheService
from src.core.crypto import CryptoService from src.core.crypto import CryptoService
from src.core.logger import logger from src.core.logger import logger
@@ -33,6 +35,36 @@ from src.services.provider_ops.types import (
BALANCE_CACHE_TTL = 86400 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: class ProviderOpsService:
""" """
Provider 操作服务 Provider 操作服务
@@ -366,6 +398,7 @@ class ProviderOpsService:
self, self,
provider_id: str, provider_id: str,
trigger_refresh: bool = True, trigger_refresh: bool = True,
allow_sync_query: bool = True,
) -> ActionResult: ) -> ActionResult:
""" """
查询余额(优先返回缓存,可触发异步刷新) 查询余额(优先返回缓存,可触发异步刷新)
@@ -373,6 +406,7 @@ class ProviderOpsService:
Args: Args:
provider_id: Provider ID provider_id: Provider ID
trigger_refresh: 是否触发后台异步刷新 trigger_refresh: 是否触发后台异步刷新
allow_sync_query: 缓存未命中时是否允许同步查询False 时仅返回缓存或触发异步刷新)
Returns: Returns:
操作结果(可能是缓存的) 操作结果(可能是缓存的)
@@ -387,19 +421,43 @@ class ProviderOpsService:
asyncio.create_task(self._refresh_balance_async(provider_id)) asyncio.create_task(self._refresh_balance_async(provider_id))
return cached return cached
# 没有缓存,同步查询一次(首次访问) # 没有缓存
logger.info(f"余额缓存未命中,同步查询: provider_id={provider_id}") if allow_sync_query:
return await self.query_balance(provider_id) # 同步查询一次(首次访问)
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: async def _refresh_balance_async(self, provider_id: str) -> None:
"""后台异步刷新余额(使用独立的数据库 session""" """
后台异步刷新余额(使用独立的数据库 session
注意:这是一个后台任务,使用独立的短生命周期 session
避免长时间占用连接池资源。
"""
db = None
try: try:
# 后台任务需要创建独立的 session因为原请求的 session 可能已关闭 # 后台任务需要创建独立的 session因为原请求的 session 可能已关闭
with create_session() as db: db = create_session()
service = ProviderOpsService(db) service = ProviderOpsService(db)
await service.query_balance(provider_id) await service.query_balance(provider_id)
except Exception as e: except Exception as e:
logger.warning(f"异步刷新余额失败: provider_id={provider_id}, error={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: async def _clear_balance_cache(self, provider_id: str) -> None:
"""清除余额缓存(认证失败时调用)""" """清除余额缓存(认证失败时调用)"""
@@ -643,6 +701,8 @@ class ProviderOpsService:
""" """
批量查询余额(优先返回缓存,后台异步刷新) 批量查询余额(优先返回缓存,后台异步刷新)
使用信号量限制并发数,避免数据库连接池耗尽。
Args: Args:
provider_ids: Provider ID 列表None 表示查询所有已配置的 provider_ids: Provider ID 列表None 表示查询所有已配置的
@@ -658,26 +718,36 @@ class ProviderOpsService:
if p.config and p.config.get("provider_ops") if p.config and p.config.get("provider_ops")
] ]
# 并行查询,使用缓存优先策略 if not provider_ids:
tasks = [ return {}
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)
results = {} # 使用信号量限制并发数,避免同时发起过多请求耗尽连接池
for provider_id, result in zip(provider_ids, results_list): concurrency = _get_batch_balance_concurrency()
if isinstance(result, Exception): semaphore = asyncio.Semaphore(concurrency)
logger.warning(f"查询余额失败: provider_id={provider_id}, error={result}") logger.debug(f"批量余额查询: providers={len(provider_ids)}, concurrency={concurrency}")
results[provider_id] = ActionResult(
status=ActionStatus.UNKNOWN_ERROR,
action_type=ProviderActionType.QUERY_BALANCE,
message=str(result),
)
else:
results[provider_id] = result
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)
# ==================== 认证验证 ==================== # ==================== 认证验证 ====================

View File

@@ -34,6 +34,7 @@ class ActionStatus(str, Enum):
"""操作执行状态""" """操作执行状态"""
SUCCESS = "success" # 成功 SUCCESS = "success" # 成功
PENDING = "pending" # 处理中(异步任务已触发,尚未完成)
AUTH_FAILED = "auth_failed" # 认证失败 AUTH_FAILED = "auth_failed" # 认证失败
AUTH_EXPIRED = "auth_expired" # 认证过期 AUTH_EXPIRED = "auth_expired" # 认证过期
RATE_LIMITED = "rate_limited" # 频率限制 RATE_LIMITED = "rate_limited" # 频率限制

View File

@@ -408,28 +408,28 @@ class UserService:
"""获取用户可用的模型 """获取用户可用的模型
通过 GlobalModel + Model 关联查询用户可用模型 通过 GlobalModel + Model 关联查询用户可用模型
逻辑:用户可用提供商 -> Provider 的 Model 实现 -> 关联的 GlobalModel 逻辑:使用 AccessRestrictions 统一处理 allowed_providers 和 allowed_models 限制
""" """
# 获取用户可用的提供商 from src.api.base.models_service import AccessRestrictions
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]
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 [] return []
# 查询这些提供商的所有活跃 Model关联 GlobalModel # 查询所有活跃 Model关联 GlobalModel
models = ( all_models = (
db.query(Model) db.query(Model)
.join(GlobalModel, Model.global_model_id == GlobalModel.id) .join(GlobalModel, Model.global_model_id == GlobalModel.id)
.filter( .filter(
and_( and_(
Model.provider_id.in_(provider_ids), Model.provider_id.in_(all_active_provider_ids),
Model.is_active == True, Model.is_active == True,
GlobalModel.is_active == True, GlobalModel.is_active == True,
) )
@@ -437,6 +437,14 @@ class UserService:
.all() .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