mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
1d9c77522a
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本，归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块，新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块，引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程，移除 base image 二级构建，统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
297 lines
9.8 KiB
Python
297 lines
9.8 KiB
Python
"""
|
||
模型可用性查询模块
|
||
|
||
将所有系统级「可用性」条件集中管理，作为模型查询的单一来源。
|
||
|
||
职责边界：
|
||
- 本模块只负责系统级可用性（对所有请求一致）
|
||
- API Key/User 的请求级访问限制由 models_service.AccessRestrictions 处理
|
||
"""
|
||
|
||
from sqlalchemy import or_, tuple_
|
||
from sqlalchemy.orm import Query, Session, contains_eager
|
||
|
||
from src.core.logger import logger
|
||
from src.models.database import (
|
||
GlobalModel,
|
||
Model,
|
||
Provider,
|
||
ProviderAPIKey,
|
||
ProviderEndpoint,
|
||
)
|
||
from src.services.provider.format import normalize_endpoint_signature
|
||
|
||
|
||
class ModelAvailabilityQuery:
|
||
"""
|
||
模型可用性查询构建器
|
||
|
||
设计原则：
|
||
1. 单一来源：所有可用性条件定义在此类中
|
||
2. 内连接 GlobalModel：未关联的 Model 不参与路由（global_model_id=NULL 不返回）
|
||
3. 完整过滤：包含 is_active 与 is_available（is_available=NULL 视为可用，兼容历史数据）
|
||
"""
|
||
|
||
@staticmethod
|
||
def base_active_models(db: Session, eager_load: bool = False) -> Query:
|
||
"""
|
||
返回基础的可用模型查询（实体为 Model）
|
||
|
||
已包含条件：
|
||
- Model.is_active = True
|
||
- Model.is_available = True 或 NULL（NULL 视为可用）
|
||
- Provider.is_active = True
|
||
- GlobalModel.is_active = True
|
||
- Model 必须关联 GlobalModel（内连接，排除 global_model_id=NULL）
|
||
|
||
Args:
|
||
db: 数据库会话
|
||
eager_load: 是否预加载 Provider 与 GlobalModel（复用 join，避免重复 JOIN）
|
||
"""
|
||
# 使用关系路径 join，与 contains_eager 兼容
|
||
query = (
|
||
db.query(Model)
|
||
.join(Model.provider)
|
||
.join(Model.global_model)
|
||
.filter(
|
||
Model.is_active.is_(True),
|
||
or_(Model.is_available.is_(True), Model.is_available.is_(None)),
|
||
Provider.is_active.is_(True),
|
||
GlobalModel.is_active.is_(True),
|
||
)
|
||
)
|
||
|
||
if eager_load:
|
||
query = query.options(
|
||
contains_eager(Model.provider),
|
||
contains_eager(Model.global_model),
|
||
)
|
||
|
||
return query
|
||
|
||
@staticmethod
|
||
def get_providers_with_active_endpoints(
|
||
db: Session,
|
||
api_formats: list[str],
|
||
) -> dict[str, set[str]]:
|
||
"""
|
||
获取有活跃端点的 Provider 及其支持的格式集合
|
||
|
||
条件：
|
||
- Provider.is_active = True（提前过滤，减少无效候选）
|
||
- ProviderEndpoint.is_active = True
|
||
- ProviderEndpoint.api_format 匹配
|
||
|
||
Returns:
|
||
{provider_id: {format1, format2, ...}}
|
||
"""
|
||
target_pairs: list[tuple[str, str]] = []
|
||
for fmt in api_formats:
|
||
if not fmt:
|
||
continue
|
||
try:
|
||
norm = normalize_endpoint_signature(fmt)
|
||
fam, kind = norm.split(":", 1)
|
||
if fam and kind:
|
||
target_pairs.append((fam, kind))
|
||
except Exception:
|
||
continue
|
||
if not target_pairs:
|
||
return {}
|
||
|
||
endpoint_rows = (
|
||
db.query(
|
||
ProviderEndpoint.provider_id,
|
||
ProviderEndpoint.api_family,
|
||
ProviderEndpoint.endpoint_kind,
|
||
)
|
||
.join(Provider, ProviderEndpoint.provider_id == Provider.id)
|
||
.filter(
|
||
Provider.is_active.is_(True),
|
||
ProviderEndpoint.is_active.is_(True),
|
||
ProviderEndpoint.api_family.isnot(None),
|
||
ProviderEndpoint.endpoint_kind.isnot(None),
|
||
tuple_(ProviderEndpoint.api_family, ProviderEndpoint.endpoint_kind).in_(
|
||
target_pairs
|
||
),
|
||
)
|
||
.all()
|
||
)
|
||
|
||
provider_to_formats: dict[str, set[str]] = {}
|
||
for provider_id, fam, kind in endpoint_rows:
|
||
if provider_id and fam and kind:
|
||
provider_to_formats.setdefault(provider_id, set()).add(
|
||
normalize_endpoint_signature(f"{fam}:{kind}")
|
||
)
|
||
|
||
return provider_to_formats
|
||
|
||
@staticmethod
|
||
def get_providers_with_active_keys(
|
||
db: Session,
|
||
provider_ids: set[str],
|
||
api_formats: list[str],
|
||
provider_to_endpoint_formats: dict[str, set[str]],
|
||
) -> set[str]:
|
||
"""
|
||
过滤出有活跃 Key 支持指定格式的 Provider
|
||
|
||
条件：
|
||
- ProviderAPIKey.is_active = True
|
||
- Key.api_formats 与 Endpoint 格式与请求格式有交集
|
||
"""
|
||
if not provider_ids:
|
||
return set()
|
||
|
||
target_formats = {normalize_endpoint_signature(f) for f in api_formats if f}
|
||
|
||
key_rows = (
|
||
db.query(ProviderAPIKey.provider_id, ProviderAPIKey.api_formats)
|
||
.filter(
|
||
ProviderAPIKey.provider_id.in_(provider_ids),
|
||
ProviderAPIKey.is_active.is_(True),
|
||
)
|
||
.all()
|
||
)
|
||
|
||
available_provider_ids: set[str] = set()
|
||
for provider_id, key_formats in key_rows:
|
||
if not provider_id:
|
||
continue
|
||
|
||
endpoint_formats = provider_to_endpoint_formats.get(provider_id)
|
||
if not endpoint_formats:
|
||
continue
|
||
|
||
# 类型兜底：key_formats 是 JSON 字段
|
||
if key_formats is None:
|
||
# None = 全支持（兼容历史数据）
|
||
key_formats_norm = set(endpoint_formats)
|
||
elif not isinstance(key_formats, list):
|
||
logger.warning(
|
||
"[ModelAvailability] Key api_formats 类型异常, provider_id={}, type={}",
|
||
provider_id,
|
||
type(key_formats).__name__,
|
||
)
|
||
continue
|
||
else:
|
||
key_formats_norm = {
|
||
normalize_endpoint_signature(str(f))
|
||
for f in key_formats
|
||
if isinstance(f, str) and f
|
||
}
|
||
|
||
if key_formats_norm & endpoint_formats & target_formats:
|
||
available_provider_ids.add(provider_id)
|
||
|
||
return available_provider_ids
|
||
|
||
@staticmethod
|
||
def get_provider_key_rules(
|
||
db: Session,
|
||
provider_ids: set[str],
|
||
api_formats: list[str],
|
||
provider_to_endpoint_formats: dict[str, set[str]],
|
||
) -> dict[str, list[tuple[list[str] | None, set[str]]]]:
|
||
"""
|
||
获取每个 Provider 的 Key 权限规则
|
||
|
||
Returns:
|
||
{provider_id: [(allowed_models, usable_formats), ...]}
|
||
|
||
注意：
|
||
- allowed_models 是 JSON 字段，此方法会进行类型兜底处理
|
||
- 非预期类型会跳过该 Key 并打日志（安全优先，不放大权限）
|
||
"""
|
||
if not provider_ids:
|
||
return {}
|
||
|
||
target_formats = {normalize_endpoint_signature(f) for f in api_formats if f}
|
||
|
||
key_rows = (
|
||
db.query(
|
||
ProviderAPIKey.id,
|
||
ProviderAPIKey.provider_id,
|
||
ProviderAPIKey.allowed_models,
|
||
ProviderAPIKey.api_formats,
|
||
)
|
||
.filter(
|
||
ProviderAPIKey.provider_id.in_(provider_ids),
|
||
ProviderAPIKey.is_active.is_(True),
|
||
)
|
||
.all()
|
||
)
|
||
|
||
provider_key_rules: dict[str, list[tuple[list[str] | None, set[str]]]] = {}
|
||
for key_id, provider_id, allowed_models_raw, key_formats in key_rows:
|
||
if not provider_id:
|
||
continue
|
||
|
||
endpoint_formats = provider_to_endpoint_formats.get(provider_id)
|
||
if not endpoint_formats:
|
||
continue
|
||
|
||
# 类型兜底：key_formats
|
||
if key_formats is None:
|
||
key_formats_norm = set(endpoint_formats)
|
||
elif not isinstance(key_formats, list):
|
||
logger.warning(
|
||
"[ModelAvailability] Key api_formats 类型异常, key_id={}, type={}",
|
||
key_id,
|
||
type(key_formats).__name__,
|
||
)
|
||
continue
|
||
else:
|
||
key_formats_norm = {
|
||
normalize_endpoint_signature(str(f))
|
||
for f in key_formats
|
||
if isinstance(f, str) and f
|
||
}
|
||
|
||
usable_formats = key_formats_norm & endpoint_formats & target_formats
|
||
if not usable_formats:
|
||
continue
|
||
|
||
# 类型兜底：allowed_models（安全优先）
|
||
allowed_models: list[str] | None
|
||
if allowed_models_raw is None:
|
||
# None = 不限制
|
||
allowed_models = None
|
||
elif isinstance(allowed_models_raw, list):
|
||
allowed_models = [m for m in allowed_models_raw if isinstance(m, str)]
|
||
else:
|
||
logger.warning(
|
||
f"[ModelAvailability] Key allowed_models 类型异常, "
|
||
f"key_id={key_id}, type={type(allowed_models_raw).__name__}, 跳过该 Key"
|
||
)
|
||
continue
|
||
|
||
provider_key_rules.setdefault(provider_id, []).append((allowed_models, usable_formats))
|
||
|
||
return provider_key_rules
|
||
|
||
@staticmethod
|
||
def find_by_global_model_name(
|
||
db: Session,
|
||
model_name: str,
|
||
provider_ids: set[str] | None = None,
|
||
eager_load: bool = False,
|
||
) -> Query:
|
||
"""
|
||
按 GlobalModel.name 查找模型
|
||
|
||
条件：
|
||
- 基础可用性条件（base_active_models）
|
||
- GlobalModel.name = model_name
|
||
- 可选：限制到指定 Provider
|
||
"""
|
||
query = ModelAvailabilityQuery.base_active_models(db, eager_load=eager_load).filter(
|
||
GlobalModel.name == model_name
|
||
)
|
||
|
||
if provider_ids is not None:
|
||
query = query.filter(Model.provider_id.in_(provider_ids))
|
||
|
||
return query
|