mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
feat: 优化批量余额查询和用户模型权限检查
- 添加批量余额查询并发限制,避免数据库连接池耗尽 - 支持余额加载 pending 状态和前端自动重试机制 - 添加用户可用模型 API,统一使用 AccessRestrictions - 修复用户表单编辑时数组引用共享导致的数据覆盖问题 - 添加数据库连接池配置说明到 .env.example
This commit is contained in:
29
.env.example
29
.env.example
@@ -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
|
||||||
|
|||||||
@@ -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')
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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 || [])]
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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 状态的 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) {
|
} 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>
|
||||||
|
|||||||
@@ -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
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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, '加载模型失败')
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|
||||||
# ==================== 认证验证 ====================
|
# ==================== 认证验证 ====================
|
||||||
|
|
||||||
|
|||||||
@@ -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" # 频率限制
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user