""" 模型可用性查询模块 将所有系统级「可用性」条件集中管理,作为模型查询的单一来源。 职责边界: - 本模块只负责系统级可用性(对所有请求一致) - 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