59 lines
2.4 KiB
Python
59 lines
2.4 KiB
Python
"""Per-account 模型访问控制(档位制)。
|
|
|
|
`users.plan` 存档位名;「档位 → 可用模型集合」定义在 `config/agent.yaml` 的 `model_tiers`。
|
|
- plan 为空 / 未知档位 → 落 `default` 档(= 基线,所有未分配用户)。
|
|
- `role == 'admin'` 始终全开,不受档位限制(管理员要能测所有模型)。
|
|
- 某档成员里出现 `"*"` → 该档全开(含未来新增模型),给内部档用。
|
|
|
|
模型 id 约定(与 list 端点 / resolve 校验一致):
|
|
- 文本模型 = `family.variant`(config/models/<family>.yaml),如 `doubao.pro`、`glm.pro52`
|
|
- 图 / 视频模型 = variant key(config/media/doubao.yaml),如 `seedream_5`、`seedance_2_fast`
|
|
两者命名不冲突(文本带点、媒体 variant 不带点),同一档集合里混放即可。
|
|
|
|
纯函数 + 读 yaml 配置,不碰 DB / HTTP —— 调用方(web 层)负责取 user 的 plan/role
|
|
并把"拒绝"翻译成 HTTP 403。这样 core 不耦合 fastapi。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from typing import Optional
|
|
|
|
DEFAULT_TIER = "default"
|
|
WILDCARD = "*"
|
|
|
|
|
|
def _tiers() -> dict[str, list[str]]:
|
|
"""从 agent.yaml 读 model_tiers;缺失 → 空 dict(→ 所有人落 default,而 default 也空 → 全禁)。
|
|
|
|
开发期不缓存,每次现读(load_config 自身轻量);改 yaml 重启 web 生效。
|
|
"""
|
|
from core.agent_builder import load_config
|
|
|
|
return load_config().get("model_tiers") or {}
|
|
|
|
|
|
def tier_name(plan: Optional[str], tiers: Optional[dict] = None) -> str:
|
|
"""plan → 实际生效的档位名;plan 为空 / 不在 tiers 里 → DEFAULT_TIER。"""
|
|
tiers = _tiers() if tiers is None else tiers
|
|
p = (plan or "").strip()
|
|
return p if p in tiers else DEFAULT_TIER
|
|
|
|
|
|
def allowed_set(plan: Optional[str], role: Optional[str]) -> Optional[set[str]]:
|
|
"""该用户可用模型 id 集合;返回 None = 全开(admin 或档位含 '*')。
|
|
|
|
None 与 空 set 语义不同:None=不设限(放行一切),空 set=一个都不许。
|
|
"""
|
|
if (role or "") == "admin":
|
|
return None
|
|
tiers = _tiers()
|
|
members = tiers.get(tier_name(plan, tiers)) or []
|
|
if WILDCARD in members:
|
|
return None
|
|
return set(members)
|
|
|
|
|
|
def is_allowed(model_id: str, plan: Optional[str], role: Optional[str]) -> bool:
|
|
"""该用户能否使用某模型 id(文本 profile 或媒体 variant)。"""
|
|
allowed = allowed_set(plan, role)
|
|
return allowed is None or model_id in allowed
|