mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
refactor: 移除 Python 后端源码,全面迁移至 Rust gateway 架构
- 删除全部 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)
This commit is contained in:
68
_deprecated_py_src/services/model/__init__.py
Normal file
68
_deprecated_py_src/services/model/__init__.py
Normal file
@@ -0,0 +1,68 @@
|
||||
"""
|
||||
模型服务模块
|
||||
|
||||
包含模型管理、成本计算等功能。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.services.model.availability import ModelAvailabilityQuery
|
||||
from src.services.model.cost import ModelCostService
|
||||
from src.services.model.fetch_scheduler import ModelFetchScheduler, get_model_fetch_scheduler
|
||||
from src.services.model.global_model import GlobalModelService
|
||||
from src.services.model.service import ModelService
|
||||
|
||||
__all__ = [
|
||||
"ModelService",
|
||||
"GlobalModelService",
|
||||
"ModelCostService",
|
||||
"ModelAvailabilityQuery",
|
||||
"ModelFetchScheduler",
|
||||
"get_model_fetch_scheduler",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
"""Lazy attribute access to avoid import-time side effects.
|
||||
|
||||
Importing `src.services.model` should not eagerly import the whole model
|
||||
service stack (scheduler/services), which can create circular imports during
|
||||
test collection.
|
||||
"""
|
||||
|
||||
if name == "ModelAvailabilityQuery":
|
||||
from src.services.model.availability import ModelAvailabilityQuery as _ModelAvailabilityQuery
|
||||
|
||||
return _ModelAvailabilityQuery
|
||||
|
||||
if name == "ModelCostService":
|
||||
from src.services.model.cost import ModelCostService as _ModelCostService
|
||||
|
||||
return _ModelCostService
|
||||
|
||||
if name == "ModelFetchScheduler":
|
||||
from src.services.model.fetch_scheduler import ModelFetchScheduler as _ModelFetchScheduler
|
||||
|
||||
return _ModelFetchScheduler
|
||||
|
||||
if name == "get_model_fetch_scheduler":
|
||||
from src.services.model.fetch_scheduler import (
|
||||
get_model_fetch_scheduler as _get_model_fetch_scheduler,
|
||||
)
|
||||
|
||||
return _get_model_fetch_scheduler
|
||||
|
||||
if name == "GlobalModelService":
|
||||
from src.services.model.global_model import GlobalModelService as _GlobalModelService
|
||||
|
||||
return _GlobalModelService
|
||||
|
||||
if name == "ModelService":
|
||||
from src.services.model.service import ModelService as _ModelService
|
||||
|
||||
return _ModelService
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
296
_deprecated_py_src/services/model/availability.py
Normal file
296
_deprecated_py_src/services/model/availability.py
Normal file
@@ -0,0 +1,296 @@
|
||||
"""
|
||||
模型可用性查询模块
|
||||
|
||||
将所有系统级「可用性」条件集中管理,作为模型查询的单一来源。
|
||||
|
||||
职责边界:
|
||||
- 本模块只负责系统级可用性(对所有请求一致)
|
||||
- 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
|
||||
1025
_deprecated_py_src/services/model/cost.py
Normal file
1025
_deprecated_py_src/services/model/cost.py
Normal file
File diff suppressed because it is too large
Load Diff
796
_deprecated_py_src/services/model/fetch_scheduler.py
Normal file
796
_deprecated_py_src/services/model/fetch_scheduler.py
Normal file
@@ -0,0 +1,796 @@
|
||||
"""
|
||||
模型自动获取调度器
|
||||
|
||||
定时从上游 API 获取可用模型列表,并更新 ProviderAPIKey 的 allowed_models。
|
||||
|
||||
功能:
|
||||
- 扫描所有启用了 auto_fetch_models 的 ProviderAPIKey
|
||||
- 调用 core.api_format 注册表获取模型列表
|
||||
- 更新 Key 的 allowed_models(保留 locked_models 中的模型)
|
||||
- 支持包含/排除规则过滤模型
|
||||
- 记录获取结果和错误信息
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import fnmatch
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import defer, joinedload, load_only
|
||||
|
||||
from src.core.cache_service import CacheService
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.logger import logger
|
||||
from src.core.provider_types import ProviderType
|
||||
from src.database import create_session
|
||||
from src.models.database import Provider, ProviderAPIKey
|
||||
from src.services.model.upstream_fetcher import (
|
||||
EndpointFetchConfig,
|
||||
UpstreamModelsFetchContext,
|
||||
build_format_to_config,
|
||||
fetch_models_for_key,
|
||||
merge_upstream_metadata,
|
||||
)
|
||||
from src.services.provider.oauth_token import resolve_oauth_access_token
|
||||
from src.services.proxy_node.resolver import resolve_effective_proxy
|
||||
from src.services.system.scheduler import get_scheduler
|
||||
|
||||
# 从环境变量读取间隔,默认 1440 分钟(1 天),限制在 60-10080 分钟之间
|
||||
_interval_env = int(os.getenv("MODEL_FETCH_INTERVAL_MINUTES", "1440"))
|
||||
MODEL_FETCH_INTERVAL_MINUTES = max(60, min(10080, _interval_env))
|
||||
|
||||
# 并发请求限制
|
||||
MAX_CONCURRENT_REQUESTS = 5
|
||||
|
||||
# 单个 Key 处理的超时时间(秒)
|
||||
KEY_FETCH_TIMEOUT_SECONDS = 120
|
||||
|
||||
# 模型获取 HTTP 请求超时时间(秒)
|
||||
# 使用较短的超时(10秒),避免不支持 /models 端点的提供商长时间阻塞
|
||||
MODEL_FETCH_HTTP_TIMEOUT = 10.0
|
||||
|
||||
# 启动时首次自动获取开关与延迟
|
||||
MODEL_FETCH_STARTUP_ENABLED = os.getenv("MODEL_FETCH_STARTUP_ENABLED", "true").lower() == "true"
|
||||
MODEL_FETCH_STARTUP_DELAY_SECONDS = max(
|
||||
0,
|
||||
int(os.getenv("MODEL_FETCH_STARTUP_DELAY_SECONDS", "10")),
|
||||
)
|
||||
|
||||
# 单批扫描的 Key 数量(仅拉取 ID,避免一次性扫描整个大号池)
|
||||
AUTO_FETCH_KEY_BATCH_SIZE = max(MAX_CONCURRENT_REQUESTS, 100)
|
||||
|
||||
# 上游模型缓存 TTL(与定时任务间隔保持一致)
|
||||
UPSTREAM_MODELS_CACHE_TTL_SECONDS = MODEL_FETCH_INTERVAL_MINUTES * 60
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PreparedModelsFetchContext:
|
||||
key_id: str
|
||||
provider_id: str
|
||||
provider_name: str
|
||||
provider_type: str
|
||||
auth_type: str
|
||||
encrypted_api_key: str
|
||||
encrypted_auth_config: str | None
|
||||
format_to_endpoint: dict[str, EndpointFetchConfig]
|
||||
proxy_config: dict[str, Any] | None
|
||||
|
||||
|
||||
def _match_pattern(model_id: str, pattern: str) -> bool:
|
||||
"""
|
||||
检查模型 ID 是否匹配模式
|
||||
|
||||
支持的通配符:
|
||||
- * 匹配任意字符(包括空)
|
||||
- ? 匹配单个字符
|
||||
|
||||
Args:
|
||||
model_id: 模型 ID
|
||||
pattern: 匹配模式
|
||||
|
||||
Returns:
|
||||
是否匹配
|
||||
"""
|
||||
return fnmatch.fnmatch(model_id.lower(), pattern.lower())
|
||||
|
||||
|
||||
def _filter_models_by_patterns(
|
||||
model_ids: set[str],
|
||||
include_patterns: list[str] | None,
|
||||
exclude_patterns: list[str] | None,
|
||||
) -> set[str]:
|
||||
"""
|
||||
根据包含/排除规则过滤模型列表
|
||||
|
||||
规则优先级:
|
||||
1. 如果 include_patterns 为空或 None,则包含所有模型
|
||||
2. 如果 include_patterns 不为空,则只包含匹配的模型
|
||||
3. exclude_patterns 总是会排除匹配的模型(优先级高于 include)
|
||||
|
||||
Args:
|
||||
model_ids: 原始模型 ID 集合
|
||||
include_patterns: 包含规则列表(支持 * 和 ? 通配符)
|
||||
exclude_patterns: 排除规则列表(支持 * 和 ? 通配符)
|
||||
|
||||
Returns:
|
||||
过滤后的模型 ID 集合
|
||||
"""
|
||||
result = set()
|
||||
|
||||
for model_id in model_ids:
|
||||
# 步骤1: 检查是否应该包含
|
||||
should_include = True
|
||||
if include_patterns:
|
||||
# 有包含规则时,必须匹配至少一个规则
|
||||
should_include = any(_match_pattern(model_id, p) for p in include_patterns)
|
||||
|
||||
if not should_include:
|
||||
continue
|
||||
|
||||
# 步骤2: 检查是否应该排除
|
||||
should_exclude = False
|
||||
if exclude_patterns:
|
||||
should_exclude = any(_match_pattern(model_id, p) for p in exclude_patterns)
|
||||
|
||||
if not should_exclude:
|
||||
result.add(model_id)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _get_upstream_models_cache_key(provider_id: str, api_key_id: str) -> str:
|
||||
"""生成上游模型缓存的 key"""
|
||||
return f"upstream_models:{provider_id}:{api_key_id}"
|
||||
|
||||
|
||||
async def get_upstream_models_from_cache(provider_id: str, api_key_id: str) -> list[dict] | None:
|
||||
"""从缓存获取上游模型列表"""
|
||||
cache_key = _get_upstream_models_cache_key(provider_id, api_key_id)
|
||||
cached = await CacheService.get(cache_key)
|
||||
if cached is not None:
|
||||
logger.debug(f"上游模型缓存命中: {cache_key}")
|
||||
return cached # type: ignore[no-any-return]
|
||||
return None
|
||||
|
||||
|
||||
async def set_upstream_models_to_cache(
|
||||
provider_id: str, api_key_id: str, models: list[dict]
|
||||
) -> None:
|
||||
"""将上游模型列表写入缓存"""
|
||||
cache_key = _get_upstream_models_cache_key(provider_id, api_key_id)
|
||||
await CacheService.set(cache_key, models, UPSTREAM_MODELS_CACHE_TTL_SECONDS)
|
||||
logger.debug(f"上游模型已缓存: {cache_key}, 数量={len(models)}")
|
||||
|
||||
|
||||
def _aggregate_models_for_cache(models: list[dict]) -> list[dict]:
|
||||
"""聚合缓存模型,按 model id 合并 api_formats,减少 Redis 占用。"""
|
||||
aggregated: dict[str, dict[str, Any]] = {}
|
||||
ordered_ids: list[str] = []
|
||||
|
||||
for model in models:
|
||||
if not isinstance(model, dict):
|
||||
continue
|
||||
|
||||
model_id = str(model.get("id") or "").strip()
|
||||
if not model_id:
|
||||
continue
|
||||
|
||||
api_format = str(model.get("api_format") or "").strip()
|
||||
existing = aggregated.get(model_id)
|
||||
|
||||
if existing is None:
|
||||
payload = dict(model)
|
||||
payload.pop("api_format", None)
|
||||
payload["api_formats"] = [api_format] if api_format else []
|
||||
aggregated[model_id] = payload
|
||||
ordered_ids.append(model_id)
|
||||
continue
|
||||
|
||||
api_formats = existing.setdefault("api_formats", [])
|
||||
if not isinstance(api_formats, list):
|
||||
api_formats = []
|
||||
existing["api_formats"] = api_formats
|
||||
|
||||
if api_format and api_format not in api_formats:
|
||||
api_formats.append(api_format)
|
||||
|
||||
for key, value in model.items():
|
||||
if key not in existing and key != "api_format":
|
||||
existing[key] = value
|
||||
|
||||
for model_id in ordered_ids:
|
||||
api_formats = aggregated[model_id].get("api_formats")
|
||||
if isinstance(api_formats, list):
|
||||
aggregated[model_id]["api_formats"] = sorted(
|
||||
{str(fmt) for fmt in api_formats if str(fmt).strip()}
|
||||
)
|
||||
|
||||
return [aggregated[model_id] for model_id in ordered_ids]
|
||||
|
||||
|
||||
async def _run_key_fetch_workers(
|
||||
key_ids: list[str],
|
||||
*,
|
||||
max_concurrent: int,
|
||||
timeout_seconds: float,
|
||||
running_predicate: Callable[[], bool],
|
||||
fetch_one: Callable[[str], Awaitable[str]],
|
||||
on_timeout: Callable[[str], None],
|
||||
on_error: Callable[[str, str], None],
|
||||
) -> tuple[int, int, int]:
|
||||
"""用固定 worker 数处理 key,避免一次性创建大量协程导致内存峰值过高。"""
|
||||
if not key_ids:
|
||||
return 0, 0, 0
|
||||
|
||||
worker_count = max(1, min(max_concurrent, len(key_ids)))
|
||||
key_queue: asyncio.Queue[str] = asyncio.Queue()
|
||||
for key_id in key_ids:
|
||||
key_queue.put_nowait(key_id)
|
||||
|
||||
async def _worker() -> tuple[int, int, int]:
|
||||
success_count = 0
|
||||
error_count = 0
|
||||
skip_count = 0
|
||||
|
||||
while True:
|
||||
if not running_predicate():
|
||||
break
|
||||
|
||||
try:
|
||||
key_id = key_queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
result = "error"
|
||||
try:
|
||||
if not running_predicate():
|
||||
result = "skip"
|
||||
else:
|
||||
result = await asyncio.wait_for(fetch_one(key_id), timeout=timeout_seconds)
|
||||
except TimeoutError:
|
||||
logger.error(f"处理 Key {key_id} 超时({timeout_seconds}s)")
|
||||
on_timeout(key_id)
|
||||
result = "error"
|
||||
except Exception as exc:
|
||||
logger.exception(f"处理 Key {key_id} 时出错")
|
||||
on_error(key_id, str(exc))
|
||||
result = "error"
|
||||
finally:
|
||||
key_queue.task_done()
|
||||
|
||||
if result == "success":
|
||||
success_count += 1
|
||||
elif result == "skip":
|
||||
skip_count += 1
|
||||
else:
|
||||
error_count += 1
|
||||
|
||||
return success_count, error_count, skip_count
|
||||
|
||||
results = await asyncio.gather(*[asyncio.create_task(_worker()) for _ in range(worker_count)])
|
||||
|
||||
success_count = sum(success for success, _, _ in results)
|
||||
error_count = sum(error for _, error, _ in results)
|
||||
skip_count = sum(skip for _, _, skip in results)
|
||||
|
||||
# 停止过程中尚未消费的队列项统一计为 skip。
|
||||
skip_count += key_queue.qsize()
|
||||
return success_count, error_count, skip_count
|
||||
|
||||
|
||||
class ModelFetchScheduler:
|
||||
"""模型自动获取调度器"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._running = False
|
||||
self._lock = asyncio.Lock()
|
||||
self._startup_task: asyncio.Task | None = None
|
||||
|
||||
async def start(self) -> None:
|
||||
"""启动调度器"""
|
||||
if self._running:
|
||||
logger.warning("ModelFetchScheduler already running")
|
||||
return
|
||||
|
||||
self._running = True
|
||||
logger.info(f"模型自动获取调度器已启动,间隔: {MODEL_FETCH_INTERVAL_MINUTES} 分钟")
|
||||
|
||||
scheduler = get_scheduler()
|
||||
scheduler.add_interval_job(
|
||||
self._scheduled_fetch_models,
|
||||
minutes=MODEL_FETCH_INTERVAL_MINUTES,
|
||||
job_id="model_auto_fetch",
|
||||
name="自动获取模型",
|
||||
)
|
||||
|
||||
# 启动时延迟执行一次,保存任务引用
|
||||
self._startup_task = asyncio.create_task(self._run_startup_task())
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""停止调度器"""
|
||||
self._running = False
|
||||
scheduler = get_scheduler()
|
||||
scheduler.remove_job("model_auto_fetch")
|
||||
|
||||
# 取消并等待启动任务完成
|
||||
if self._startup_task and not self._startup_task.done():
|
||||
self._startup_task.cancel()
|
||||
try:
|
||||
await self._startup_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
logger.info("模型自动获取调度器已停止")
|
||||
|
||||
async def _run_startup_task(self) -> None:
|
||||
"""启动时执行的初始化任务"""
|
||||
try:
|
||||
if not MODEL_FETCH_STARTUP_ENABLED:
|
||||
logger.info("启动时模型自动获取已禁用(MODEL_FETCH_STARTUP_ENABLED=false)")
|
||||
return
|
||||
|
||||
if MODEL_FETCH_STARTUP_DELAY_SECONDS > 0:
|
||||
await asyncio.sleep(MODEL_FETCH_STARTUP_DELAY_SECONDS)
|
||||
|
||||
if not self._running:
|
||||
return
|
||||
logger.info("启动时执行首次模型获取...")
|
||||
await self._perform_fetch_all_keys()
|
||||
except asyncio.CancelledError:
|
||||
logger.debug("启动任务被取消")
|
||||
raise
|
||||
except Exception:
|
||||
logger.exception("启动时模型获取出错")
|
||||
|
||||
async def _scheduled_fetch_models(self) -> None:
|
||||
"""定时任务入口"""
|
||||
if not self._running:
|
||||
return
|
||||
async with self._lock:
|
||||
await self._perform_fetch_all_keys()
|
||||
|
||||
def _list_auto_fetch_key_id_batch(
|
||||
self,
|
||||
*,
|
||||
after_id: str | None = None,
|
||||
limit: int = AUTO_FETCH_KEY_BATCH_SIZE,
|
||||
) -> list[str]:
|
||||
"""分批返回启用自动获取模型的 Key ID。"""
|
||||
with create_session() as db:
|
||||
query = db.query(ProviderAPIKey.id).filter(
|
||||
ProviderAPIKey.auto_fetch_models == True, # noqa: E712
|
||||
ProviderAPIKey.is_active == True, # noqa: E712
|
||||
)
|
||||
if after_id:
|
||||
query = query.filter(ProviderAPIKey.id > after_id)
|
||||
return [row[0] for row in query.order_by(ProviderAPIKey.id.asc()).limit(limit).all()]
|
||||
|
||||
async def _perform_fetch_all_keys(self) -> None:
|
||||
"""获取所有启用自动获取的 Key,并以固定批次/并发节奏拉取。"""
|
||||
logger.info("开始自动获取模型任务...")
|
||||
|
||||
success_count = 0
|
||||
error_count = 0
|
||||
skip_count = 0
|
||||
total_count = 0
|
||||
last_id: str | None = None
|
||||
batch_index = 0
|
||||
|
||||
while self._running:
|
||||
key_ids = self._list_auto_fetch_key_id_batch(after_id=last_id)
|
||||
if not key_ids:
|
||||
break
|
||||
|
||||
batch_index += 1
|
||||
total_count += len(key_ids)
|
||||
last_id = key_ids[-1]
|
||||
logger.info(
|
||||
"自动获取模型任务处理第 {} 批 Key: {} 个",
|
||||
batch_index,
|
||||
len(key_ids),
|
||||
)
|
||||
|
||||
batch_success, batch_error, batch_skip = await _run_key_fetch_workers(
|
||||
key_ids,
|
||||
max_concurrent=MAX_CONCURRENT_REQUESTS,
|
||||
timeout_seconds=KEY_FETCH_TIMEOUT_SECONDS,
|
||||
running_predicate=lambda: self._running,
|
||||
fetch_one=self._fetch_models_for_key_by_id,
|
||||
on_timeout=lambda key_id: self._update_key_error(
|
||||
key_id, f"Timeout after {KEY_FETCH_TIMEOUT_SECONDS}s"
|
||||
),
|
||||
on_error=self._update_key_error,
|
||||
)
|
||||
success_count += batch_success
|
||||
error_count += batch_error
|
||||
skip_count += batch_skip
|
||||
|
||||
if len(key_ids) < AUTO_FETCH_KEY_BATCH_SIZE:
|
||||
break
|
||||
|
||||
await asyncio.sleep(0)
|
||||
|
||||
if total_count == 0:
|
||||
logger.debug("没有启用自动获取模型的 Key")
|
||||
return
|
||||
|
||||
logger.info(
|
||||
"自动获取模型任务完成: 总数={}, 成功={}, 失败={}, 跳过={}",
|
||||
total_count,
|
||||
success_count,
|
||||
error_count,
|
||||
skip_count,
|
||||
)
|
||||
|
||||
def _update_key_error(self, key_id: str, error_msg: str) -> None:
|
||||
"""更新 Key 的错误信息(独立事务)"""
|
||||
try:
|
||||
with create_session() as db:
|
||||
key = (
|
||||
db.query(ProviderAPIKey)
|
||||
.options(
|
||||
defer(ProviderAPIKey.adjustment_history),
|
||||
defer(ProviderAPIKey.utilization_samples),
|
||||
defer(ProviderAPIKey.health_by_format),
|
||||
defer(ProviderAPIKey.circuit_breaker_by_format),
|
||||
defer(ProviderAPIKey.upstream_metadata),
|
||||
defer(ProviderAPIKey.allowed_models),
|
||||
)
|
||||
.filter(ProviderAPIKey.id == key_id)
|
||||
.first()
|
||||
)
|
||||
if key:
|
||||
key.last_models_fetch_at = datetime.now(timezone.utc)
|
||||
key.last_models_fetch_error = error_msg
|
||||
db.commit()
|
||||
except Exception:
|
||||
logger.exception(f"更新 Key {key_id} 错误信息失败")
|
||||
|
||||
async def _fetch_models_for_key_by_id(self, key_id: str) -> str:
|
||||
"""
|
||||
根据 Key ID 获取模型并更新,返回结果状态
|
||||
|
||||
优化:分两个阶段处理,HTTP 请求期间不持有数据库连接,避免阻塞其他请求
|
||||
"""
|
||||
# ========== 阶段 1:准备数据(短暂持有连接)==========
|
||||
prepared = self._prepare_fetch_context(key_id)
|
||||
if prepared is None:
|
||||
return "skip"
|
||||
if isinstance(prepared, str):
|
||||
return prepared # "error" or "skip"
|
||||
|
||||
# Resolve auth (incl. lazy OAuth refresh) without holding a DB session.
|
||||
api_key_value: str = ""
|
||||
auth_config: dict[str, Any] | None = None
|
||||
|
||||
if prepared.auth_type == "oauth":
|
||||
# Use request_builder's lazy refresh logic and persist refreshed token back to DB.
|
||||
# Endpoint signature is only used for tracing/debug; auth logic doesn't depend on it.
|
||||
endpoint_api_format = (
|
||||
"gemini:chat"
|
||||
if prepared.provider_type.lower() == ProviderType.ANTIGRAVITY
|
||||
else None
|
||||
)
|
||||
try:
|
||||
resolved = await resolve_oauth_access_token(
|
||||
key_id=prepared.key_id,
|
||||
encrypted_api_key=prepared.encrypted_api_key,
|
||||
encrypted_auth_config=prepared.encrypted_auth_config,
|
||||
provider_proxy_config=prepared.proxy_config,
|
||||
endpoint_api_format=endpoint_api_format,
|
||||
)
|
||||
api_key_value = resolved.access_token
|
||||
auth_config = resolved.decrypted_auth_config
|
||||
except Exception as e:
|
||||
self._update_key_error(prepared.key_id, f"OAuth token resolution failed: {e}")
|
||||
return "error"
|
||||
else:
|
||||
is_vertex_service_account = (
|
||||
prepared.provider_type.lower() == ProviderType.VERTEX_AI.value
|
||||
and prepared.auth_type in ("service_account", "vertex_ai")
|
||||
)
|
||||
|
||||
if is_vertex_service_account:
|
||||
api_key_value = "__placeholder__"
|
||||
else:
|
||||
try:
|
||||
api_key_value = crypto_service.decrypt(prepared.encrypted_api_key)
|
||||
except Exception:
|
||||
self._update_key_error(prepared.key_id, "Decrypt error")
|
||||
return "error"
|
||||
|
||||
# Best-effort: decrypt auth_config if present (e.g. Antigravity project_id / Vertex SA JSON).
|
||||
if prepared.encrypted_auth_config:
|
||||
try:
|
||||
parsed = json.loads(crypto_service.decrypt(prepared.encrypted_auth_config))
|
||||
auth_config = parsed if isinstance(parsed, dict) else None
|
||||
except Exception:
|
||||
auth_config = None
|
||||
|
||||
fetch_ctx = UpstreamModelsFetchContext(
|
||||
provider_type=prepared.provider_type,
|
||||
api_key_value=api_key_value,
|
||||
format_to_endpoint=prepared.format_to_endpoint,
|
||||
proxy_config=prepared.proxy_config,
|
||||
auth_config=auth_config,
|
||||
)
|
||||
|
||||
# ========== 阶段 2:HTTP 请求(不持有数据库连接)==========
|
||||
# 使用较短的超时时间(10秒),避免长时间阻塞
|
||||
all_models, errors, has_success, upstream_metadata = await fetch_models_for_key(
|
||||
fetch_ctx,
|
||||
timeout_seconds=MODEL_FETCH_HTTP_TIMEOUT,
|
||||
)
|
||||
|
||||
# ========== 阶段 3:更新数据库(获取新连接)==========
|
||||
return await self._update_key_after_fetch(
|
||||
key_id=prepared.key_id,
|
||||
provider_id=prepared.provider_id,
|
||||
provider_name=prepared.provider_name,
|
||||
all_models=all_models,
|
||||
errors=errors,
|
||||
has_success=has_success,
|
||||
upstream_metadata=upstream_metadata,
|
||||
)
|
||||
|
||||
def _prepare_fetch_context(self, key_id: str) -> PreparedModelsFetchContext | str | None:
|
||||
"""
|
||||
准备获取模型所需的上下文数据
|
||||
|
||||
Returns:
|
||||
- PreparedModelsFetchContext: 准备好的上下文(不包含解密后的 token)
|
||||
- "skip": 跳过该 Key
|
||||
- "error": 出错
|
||||
- None: Key 不存在
|
||||
"""
|
||||
with create_session() as db:
|
||||
key = (
|
||||
db.query(ProviderAPIKey)
|
||||
.options(
|
||||
load_only(
|
||||
ProviderAPIKey.id,
|
||||
ProviderAPIKey.is_active,
|
||||
ProviderAPIKey.auto_fetch_models,
|
||||
ProviderAPIKey.api_key,
|
||||
ProviderAPIKey.auth_type,
|
||||
ProviderAPIKey.auth_config,
|
||||
ProviderAPIKey.provider_id,
|
||||
ProviderAPIKey.proxy,
|
||||
),
|
||||
)
|
||||
.filter(ProviderAPIKey.id == key_id)
|
||||
.first()
|
||||
)
|
||||
|
||||
if not key:
|
||||
logger.warning(f"Key {key_id} 不存在,跳过")
|
||||
return None
|
||||
|
||||
if not key.is_active or not key.auto_fetch_models:
|
||||
logger.debug(f"Key {key_id} 已禁用或关闭自动获取,跳过")
|
||||
return "skip"
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
provider_id = key.provider_id
|
||||
|
||||
# 获取 Provider 和 Endpoints
|
||||
provider = (
|
||||
db.query(Provider)
|
||||
.options(joinedload(Provider.endpoints))
|
||||
.filter(Provider.id == provider_id)
|
||||
.first()
|
||||
)
|
||||
|
||||
if not provider:
|
||||
logger.warning(f"Provider {provider_id} 不存在,跳过 Key {key.id}")
|
||||
key.last_models_fetch_error = "Provider not found"
|
||||
key.last_models_fetch_at = now
|
||||
db.commit()
|
||||
return "error"
|
||||
|
||||
auth_type = getattr(key, "auth_type", "api_key") or "api_key"
|
||||
provider_type = str(getattr(provider, "provider_type", "") or "")
|
||||
is_vertex_service_account = (
|
||||
provider_type.strip().lower() == ProviderType.VERTEX_AI.value
|
||||
and auth_type in ("service_account", "vertex_ai")
|
||||
)
|
||||
if auth_type in ("service_account", "vertex_ai") and not is_vertex_service_account:
|
||||
key.last_models_fetch_error = (
|
||||
"auto_fetch_models 暂不支持 Service Account 类型的 Key"
|
||||
)
|
||||
key.last_models_fetch_at = now
|
||||
db.commit()
|
||||
logger.info(f"Key {key.id} 为 Service Account 类型,跳过自动获取模型")
|
||||
return "skip"
|
||||
|
||||
# 基础校验:必须有 api_key(OAuth: 加密 access_token;API Key: 加密 key)
|
||||
if not key.api_key:
|
||||
logger.warning(f"Key {key.id} 没有 API Key,跳过")
|
||||
key.last_models_fetch_error = "No API key configured"
|
||||
key.last_models_fetch_at = now
|
||||
db.commit()
|
||||
return "error"
|
||||
|
||||
# 构建 api_format -> EndpointFetchConfig 映射(纯数据,session 无关)
|
||||
format_to_endpoint = build_format_to_config(provider.endpoints) # type: ignore[attr-defined]
|
||||
|
||||
if not format_to_endpoint:
|
||||
logger.warning(f"Provider {provider.name} 没有活跃的端点,跳过 Key {key.id}")
|
||||
key.last_models_fetch_error = "No active endpoints"
|
||||
key.last_models_fetch_at = now
|
||||
db.commit()
|
||||
return "error"
|
||||
encrypted_auth_config = getattr(key, "auth_config", None)
|
||||
|
||||
return PreparedModelsFetchContext(
|
||||
key_id=key_id,
|
||||
provider_id=provider_id,
|
||||
provider_name=provider.name,
|
||||
provider_type=provider_type,
|
||||
auth_type=auth_type,
|
||||
encrypted_api_key=str(key.api_key),
|
||||
encrypted_auth_config=(
|
||||
encrypted_auth_config if isinstance(encrypted_auth_config, str) else None
|
||||
),
|
||||
format_to_endpoint=format_to_endpoint,
|
||||
proxy_config=resolve_effective_proxy(
|
||||
getattr(provider, "proxy", None), getattr(key, "proxy", None)
|
||||
),
|
||||
)
|
||||
|
||||
async def _update_key_after_fetch(
|
||||
self,
|
||||
key_id: str,
|
||||
provider_id: str,
|
||||
provider_name: str,
|
||||
all_models: list[dict],
|
||||
errors: list[str],
|
||||
has_success: bool,
|
||||
upstream_metadata: dict[str, Any] | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
HTTP 请求完成后更新数据库
|
||||
|
||||
使用新的数据库连接来更新 Key 的 allowed_models
|
||||
"""
|
||||
now = datetime.now(timezone.utc)
|
||||
|
||||
with create_session() as db:
|
||||
# 重新获取 Key(因为之前的连接已关闭)
|
||||
# defer 不需要的大 JSON 字段,减少内存占用
|
||||
key = (
|
||||
db.query(ProviderAPIKey)
|
||||
.options(
|
||||
defer(ProviderAPIKey.adjustment_history),
|
||||
defer(ProviderAPIKey.utilization_samples),
|
||||
defer(ProviderAPIKey.health_by_format),
|
||||
defer(ProviderAPIKey.circuit_breaker_by_format),
|
||||
)
|
||||
.filter(ProviderAPIKey.id == key_id)
|
||||
.first()
|
||||
)
|
||||
if not key:
|
||||
logger.warning(f"Key {key_id} 在更新时不存在")
|
||||
return "error"
|
||||
|
||||
# 记录获取时间
|
||||
key.last_models_fetch_at = now
|
||||
|
||||
# 如果没有任何成功的响应,不更新 allowed_models(保留旧数据)
|
||||
if not has_success:
|
||||
error_msg = "; ".join(errors) if errors else "All endpoints failed"
|
||||
key.last_models_fetch_error = error_msg
|
||||
logger.warning(
|
||||
f"Provider {provider_name} Key {key.id} 所有端点获取失败,保留现有模型列表"
|
||||
)
|
||||
db.commit()
|
||||
return "error"
|
||||
|
||||
# 有成功的响应,清除错误状态
|
||||
key.last_models_fetch_error = None
|
||||
|
||||
# 最佳努力:保存上游元数据(如 Antigravity 配额信息)
|
||||
if upstream_metadata and isinstance(upstream_metadata, dict):
|
||||
key.upstream_metadata = merge_upstream_metadata(
|
||||
key.upstream_metadata, upstream_metadata
|
||||
)
|
||||
|
||||
# 去重获取模型 ID 列表
|
||||
fetched_model_ids: set[str] = set()
|
||||
for model in all_models:
|
||||
model_id = model.get("id")
|
||||
if model_id:
|
||||
fetched_model_ids.add(model_id)
|
||||
|
||||
logger.info(
|
||||
f"Provider {provider_name} Key {key.id} 获取到 {len(fetched_model_ids)} 个唯一模型"
|
||||
)
|
||||
|
||||
# 写入上游模型缓存(按 model id 聚合 api_formats,减少 Redis 内存占用)
|
||||
unique_models = _aggregate_models_for_cache(all_models)
|
||||
await set_upstream_models_to_cache(provider_id, key.id, unique_models)
|
||||
|
||||
# 更新 allowed_models(保留 locked_models)
|
||||
has_changed = self._update_key_allowed_models(key, fetched_model_ids)
|
||||
|
||||
db.commit()
|
||||
|
||||
# 如果白名单有变化,触发缓存失效和自动关联检查
|
||||
if has_changed and provider_id:
|
||||
from src.services.model.global_model import on_key_allowed_models_changed
|
||||
|
||||
# 使用新会话处理后续操作
|
||||
with create_session() as db2:
|
||||
await on_key_allowed_models_changed(
|
||||
db=db2,
|
||||
provider_id=provider_id,
|
||||
allowed_models=list(key.allowed_models or []),
|
||||
)
|
||||
|
||||
return "success"
|
||||
|
||||
def _update_key_allowed_models(self, key: ProviderAPIKey, fetched_model_ids: set[str]) -> bool:
|
||||
"""
|
||||
更新 Key 的 allowed_models,保留 locked_models,应用过滤规则
|
||||
|
||||
Returns:
|
||||
bool: 是否有变化
|
||||
"""
|
||||
# 获取当前锁定的模型
|
||||
locked_models = set(key.locked_models or [])
|
||||
|
||||
# 应用包含/排除过滤规则
|
||||
include_patterns = key.model_include_patterns
|
||||
exclude_patterns = key.model_exclude_patterns
|
||||
|
||||
filtered_model_ids = _filter_models_by_patterns(
|
||||
fetched_model_ids, include_patterns, exclude_patterns
|
||||
)
|
||||
|
||||
# 记录过滤结果
|
||||
if include_patterns or exclude_patterns:
|
||||
filtered_count = len(fetched_model_ids) - len(filtered_model_ids)
|
||||
if filtered_count > 0:
|
||||
logger.info(
|
||||
f"Key {key.id} 过滤规则生效: 原始 {len(fetched_model_ids)} 个模型, "
|
||||
f"过滤后 {len(filtered_model_ids)} 个 (排除 {filtered_count} 个)"
|
||||
)
|
||||
|
||||
# 新的 allowed_models = 过滤后的模型 + 锁定的模型
|
||||
# 锁定模型无论上游是否返回都会保留
|
||||
new_allowed_models = list(filtered_model_ids | locked_models)
|
||||
new_allowed_models.sort() # 保持顺序稳定
|
||||
|
||||
# 检查是否有变化
|
||||
current_allowed = set(key.allowed_models or [])
|
||||
new_allowed_set = set(new_allowed_models)
|
||||
|
||||
if current_allowed != new_allowed_set:
|
||||
added = new_allowed_set - current_allowed
|
||||
removed = current_allowed - new_allowed_set
|
||||
if added:
|
||||
logger.info(f"Key {key.id} 新增模型: {sorted(added)}")
|
||||
if removed:
|
||||
logger.info(f"Key {key.id} 移除模型: {sorted(removed)}")
|
||||
|
||||
key.allowed_models = new_allowed_models
|
||||
return True
|
||||
else:
|
||||
logger.debug(f"Key {key.id} 模型列表无变化")
|
||||
return False
|
||||
|
||||
|
||||
# 单例模式
|
||||
_model_fetch_scheduler: ModelFetchScheduler | None = None
|
||||
|
||||
|
||||
def get_model_fetch_scheduler() -> ModelFetchScheduler:
|
||||
"""获取模型获取调度器单例"""
|
||||
global _model_fetch_scheduler
|
||||
if _model_fetch_scheduler is None:
|
||||
_model_fetch_scheduler = ModelFetchScheduler()
|
||||
return _model_fetch_scheduler
|
||||
625
_deprecated_py_src/services/model/global_model.py
Normal file
625
_deprecated_py_src/services/model/global_model.py
Normal file
@@ -0,0 +1,625 @@
|
||||
"""
|
||||
GlobalModel 服务层
|
||||
|
||||
提供 GlobalModel 的 CRUD 操作、查询和统计功能
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import cast
|
||||
|
||||
from sqlalchemy import delete as sa_delete
|
||||
from sqlalchemy import func
|
||||
from sqlalchemy.orm import Session, joinedload, load_only
|
||||
|
||||
from src.core.exceptions import InvalidRequestException, NotFoundException
|
||||
from src.core.logger import logger
|
||||
from src.models.database import GlobalModel, Model
|
||||
from src.models.pydantic_models import GlobalModelUpdate
|
||||
|
||||
|
||||
async def on_key_allowed_models_changed(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
allowed_models: list[str] | None = None,
|
||||
skip_disassociate: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Key 的 allowed_models 变更后的统一处理
|
||||
|
||||
包括:
|
||||
1. 触发缓存失效(包括 /v1/models 列表缓存)
|
||||
2. 检查并自动关联匹配的 GlobalModel(仅当提供 allowed_models 时)
|
||||
3. 检查并自动解除不再匹配的 GlobalModel 关联(可通过 skip_disassociate 跳过)
|
||||
|
||||
Args:
|
||||
db: 数据库 Session
|
||||
provider_id: Provider ID
|
||||
allowed_models: 更新后的 allowed_models 列表
|
||||
- 提供非空列表:触发自动关联和解除关联检查
|
||||
- 提供空列表或 None:仅触发解除关联检查(用于 Key 删除场景)
|
||||
skip_disassociate: 是否跳过解除关联检查
|
||||
- True:跳过(用于删除 allowed_models 为 null 的 Key 时)
|
||||
- False:执行检查(默认)
|
||||
"""
|
||||
from src.services.cache.invalidation import get_cache_invalidation_service
|
||||
|
||||
# 1. 触发缓存失效
|
||||
cache_service = get_cache_invalidation_service()
|
||||
await cache_service.on_key_allowed_models_changed(provider_id)
|
||||
|
||||
# 2. 检查并自动关联 GlobalModel(仅当提供非空 allowed_models 时)
|
||||
if allowed_models:
|
||||
GlobalModelService.auto_associate_provider_by_key_whitelist(
|
||||
db=db,
|
||||
provider_id=provider_id,
|
||||
allowed_models=allowed_models,
|
||||
)
|
||||
|
||||
# 3. 检查并自动解除不再匹配的 GlobalModel 关联
|
||||
if not skip_disassociate:
|
||||
GlobalModelService.auto_disassociate_provider_by_key_whitelist(
|
||||
db=db,
|
||||
provider_id=provider_id,
|
||||
)
|
||||
|
||||
|
||||
class GlobalModelService:
|
||||
"""GlobalModel 服务"""
|
||||
|
||||
@staticmethod
|
||||
def get_global_model(db: Session, global_model_id: str) -> GlobalModel:
|
||||
"""
|
||||
获取单个 GlobalModel
|
||||
|
||||
Args:
|
||||
global_model_id: GlobalModel 的 UUID 或 name
|
||||
"""
|
||||
# 先尝试通过 ID 查找
|
||||
global_model = db.query(GlobalModel).filter(GlobalModel.id == global_model_id).first()
|
||||
|
||||
# 如果没找到,尝试通过 name 查找
|
||||
if not global_model:
|
||||
global_model = db.query(GlobalModel).filter(GlobalModel.name == global_model_id).first()
|
||||
|
||||
if not global_model:
|
||||
raise NotFoundException(f"GlobalModel {global_model_id} not found")
|
||||
return global_model
|
||||
|
||||
@staticmethod
|
||||
def get_global_model_by_name(db: Session, name: str) -> GlobalModel | None:
|
||||
"""通过名称获取 GlobalModel"""
|
||||
return db.query(GlobalModel).filter(GlobalModel.name == name).first()
|
||||
|
||||
@staticmethod
|
||||
def list_global_models(
|
||||
db: Session,
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
is_active: bool | None = None,
|
||||
search: str | None = None,
|
||||
) -> list[GlobalModel]:
|
||||
"""列出 GlobalModel"""
|
||||
query = db.query(GlobalModel)
|
||||
|
||||
if is_active is not None:
|
||||
query = query.filter(GlobalModel.is_active == is_active)
|
||||
|
||||
if search:
|
||||
search_pattern = f"%{search}%"
|
||||
query = query.filter(
|
||||
(GlobalModel.name.ilike(search_pattern))
|
||||
| (GlobalModel.display_name.ilike(search_pattern))
|
||||
)
|
||||
|
||||
# 按名称排序
|
||||
query = query.order_by(GlobalModel.name)
|
||||
|
||||
return query.offset(skip).limit(limit).all()
|
||||
|
||||
@staticmethod
|
||||
def create_global_model(
|
||||
db: Session,
|
||||
name: str,
|
||||
display_name: str,
|
||||
is_active: bool | None = True,
|
||||
# 按次计费配置
|
||||
default_price_per_request: float | None = None,
|
||||
# 阶梯计费配置(必填)
|
||||
default_tiered_pricing: dict | None = None,
|
||||
# Key 能力配置
|
||||
supported_capabilities: list[str] | None = None,
|
||||
# 模型配置(JSON)
|
||||
config: dict | None = None,
|
||||
) -> GlobalModel:
|
||||
"""创建 GlobalModel"""
|
||||
# 检查名称是否已存在
|
||||
existing = GlobalModelService.get_global_model_by_name(db, name)
|
||||
if existing:
|
||||
raise InvalidRequestException(f"GlobalModel with name '{name}' already exists")
|
||||
|
||||
global_model = GlobalModel(
|
||||
name=name,
|
||||
display_name=display_name,
|
||||
is_active=is_active,
|
||||
# 按次计费配置
|
||||
default_price_per_request=default_price_per_request,
|
||||
# 阶梯计费配置
|
||||
default_tiered_pricing=default_tiered_pricing,
|
||||
# Key 能力配置
|
||||
supported_capabilities=supported_capabilities,
|
||||
# 模型配置(JSON)
|
||||
config=config,
|
||||
)
|
||||
|
||||
db.add(global_model)
|
||||
db.commit()
|
||||
db.refresh(global_model)
|
||||
|
||||
return global_model
|
||||
|
||||
@staticmethod
|
||||
def update_global_model(
|
||||
db: Session,
|
||||
global_model_id: str,
|
||||
update_data: GlobalModelUpdate,
|
||||
) -> GlobalModel:
|
||||
"""
|
||||
更新 GlobalModel
|
||||
|
||||
使用 exclude_unset=True 来区分"未提供字段"和"显式设置为 None":
|
||||
- 未提供的字段不会被更新
|
||||
- 显式设置为 None 的字段会被更新为 None(置空)
|
||||
"""
|
||||
global_model = GlobalModelService.get_global_model(db, global_model_id)
|
||||
|
||||
# 只更新显式设置的字段(包括显式设置为 None 的情况)
|
||||
data_dict = update_data.model_dump(exclude_unset=True)
|
||||
|
||||
# 处理阶梯计费配置:如果是 TieredPricingConfig 对象,转换为 dict
|
||||
if "default_tiered_pricing" in data_dict:
|
||||
tiered_pricing = data_dict["default_tiered_pricing"]
|
||||
if tiered_pricing is not None and hasattr(tiered_pricing, "model_dump"):
|
||||
data_dict["default_tiered_pricing"] = tiered_pricing.model_dump()
|
||||
|
||||
for field, value in data_dict.items():
|
||||
setattr(global_model, field, value)
|
||||
|
||||
db.commit()
|
||||
db.refresh(global_model)
|
||||
|
||||
return global_model
|
||||
|
||||
@staticmethod
|
||||
def delete_global_model(db: Session, global_model_id: str) -> None:
|
||||
"""
|
||||
删除 GlobalModel
|
||||
|
||||
默认行为: 级联删除所有关联的 Provider 模型实现
|
||||
注意: 不清理 API Key 和 User 的 allowed_models 引用,
|
||||
保留无效引用可让用户在前端看到"已失效"的模型,便于手动清理或等待重建同名模型
|
||||
"""
|
||||
global_model = GlobalModelService.get_global_model(db, global_model_id)
|
||||
|
||||
# 批量删除所有关联的 Provider 模型实现
|
||||
assoc_count = (
|
||||
db.query(func.count(Model.id)).filter(Model.global_model_id == global_model.id).scalar()
|
||||
)
|
||||
if assoc_count:
|
||||
logger.info(
|
||||
f"删除 GlobalModel {global_model.name} 的 {assoc_count} 个关联 Provider 模型"
|
||||
)
|
||||
db.execute(sa_delete(Model).where(Model.global_model_id == global_model.id))
|
||||
|
||||
# 删除 GlobalModel
|
||||
db.delete(global_model)
|
||||
db.commit()
|
||||
|
||||
@staticmethod
|
||||
def get_global_model_stats(db: Session, global_model_id: str) -> dict:
|
||||
"""获取 GlobalModel 统计信息"""
|
||||
global_model = GlobalModelService.get_global_model(db, global_model_id)
|
||||
|
||||
# 统计关联的 Model 数量(使用 global_model.id,预加载 provider 关联)
|
||||
models = (
|
||||
db.query(Model)
|
||||
.options(joinedload(Model.provider))
|
||||
.filter(Model.global_model_id == global_model.id)
|
||||
.all()
|
||||
)
|
||||
|
||||
# 统计支持的 Provider 数量
|
||||
provider_ids = {model.provider_id for model in models}
|
||||
|
||||
# 从阶梯计费中提取价格范围
|
||||
input_prices = []
|
||||
output_prices = []
|
||||
for m in models:
|
||||
tiered = m.get_effective_tiered_pricing()
|
||||
if tiered and tiered.get("tiers"):
|
||||
first_tier = tiered["tiers"][0]
|
||||
if first_tier.get("input_price_per_1m") is not None:
|
||||
input_prices.append(first_tier["input_price_per_1m"])
|
||||
if first_tier.get("output_price_per_1m") is not None:
|
||||
output_prices.append(first_tier["output_price_per_1m"])
|
||||
|
||||
return {
|
||||
"global_model_id": global_model.id,
|
||||
"name": global_model.name,
|
||||
"total_models": len(models),
|
||||
"total_providers": len(provider_ids),
|
||||
"price_range": {
|
||||
"min_input": min(input_prices) if input_prices else None,
|
||||
"max_input": max(input_prices) if input_prices else None,
|
||||
"min_output": min(output_prices) if output_prices else None,
|
||||
"max_output": max(output_prices) if output_prices else None,
|
||||
},
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def batch_assign_to_providers(
|
||||
db: Session,
|
||||
global_model_id: str,
|
||||
provider_ids: list[str],
|
||||
create_models: bool = False,
|
||||
) -> dict:
|
||||
"""批量为多个 Provider 添加 GlobalModel 实现"""
|
||||
|
||||
global_model = GlobalModelService.get_global_model(db, global_model_id)
|
||||
|
||||
results = {
|
||||
"success": [],
|
||||
"errors": [],
|
||||
}
|
||||
|
||||
for provider_id in provider_ids:
|
||||
try:
|
||||
# 检查该 Provider 是否已有该 GlobalModel 的实现(使用 global_model.id)
|
||||
existing_model = (
|
||||
db.query(Model)
|
||||
.filter(
|
||||
Model.provider_id == provider_id,
|
||||
Model.global_model_id == global_model.id,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
if existing_model:
|
||||
results["errors"].append(
|
||||
{
|
||||
"provider_id": provider_id,
|
||||
"error": "Model already exists for this provider",
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
if create_models:
|
||||
# 创建新的 Model(价格和能力设为 None,继承 GlobalModel 默认值)
|
||||
model = Model(
|
||||
provider_id=provider_id,
|
||||
global_model_id=global_model.id,
|
||||
provider_model_name=global_model.name, # 默认使用 GlobalModel name
|
||||
# 计费设为 None,使用 GlobalModel 默认值
|
||||
price_per_request=None,
|
||||
tiered_pricing=None,
|
||||
# 能力设为 None,使用 GlobalModel 默认值
|
||||
supports_vision=None,
|
||||
supports_function_calling=None,
|
||||
supports_streaming=None,
|
||||
supports_extended_thinking=None,
|
||||
is_active=True,
|
||||
)
|
||||
db.add(model)
|
||||
db.commit()
|
||||
|
||||
results["success"].append(
|
||||
{"provider_id": provider_id, "model_id": model.id, "created": True}
|
||||
)
|
||||
else:
|
||||
results["errors"].append(
|
||||
{
|
||||
"provider_id": provider_id,
|
||||
"error": "create_models=False, no existing model found",
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
results["errors"].append({"provider_id": provider_id, "error": str(e)})
|
||||
|
||||
db.commit()
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def auto_associate_provider_by_key_whitelist(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
allowed_models: list[str],
|
||||
) -> dict:
|
||||
"""
|
||||
根据 Key 白名单自动关联 Provider 到匹配的 GlobalModel
|
||||
|
||||
当 Key 的 allowed_models 更新后调用此方法,检查所有 GlobalModel 的映射规则,
|
||||
如果有映射规则匹配到 Key 白名单中的模型,且 Provider 尚未关联到该 GlobalModel,
|
||||
则自动创建关联。
|
||||
|
||||
Args:
|
||||
db: 数据库 Session
|
||||
provider_id: Provider ID
|
||||
allowed_models: Key 的白名单模型列表
|
||||
|
||||
Returns:
|
||||
Dict: 包含 success 和 errors 列表
|
||||
"""
|
||||
from src.core.model_permissions import match_model_with_pattern
|
||||
from src.models.database import Provider
|
||||
|
||||
results: dict[str, list[dict]] = {
|
||||
"success": [],
|
||||
"errors": [],
|
||||
}
|
||||
|
||||
if not allowed_models:
|
||||
return results
|
||||
|
||||
# 获取 Provider
|
||||
provider = db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
if not provider:
|
||||
logger.warning(f"Provider {provider_id} not found for auto-association")
|
||||
return results
|
||||
|
||||
# 获取该 Provider 已关联的 GlobalModel ID 集合
|
||||
existing_associations = (
|
||||
db.query(Model.global_model_id, Model.provider_model_name)
|
||||
.filter(Model.provider_id == provider_id)
|
||||
.all()
|
||||
)
|
||||
linked_global_model_ids: set[str] = {row[0] for row in existing_associations if row[0]}
|
||||
# 同时获取已存在的 provider_model_name 集合,避免唯一约束冲突
|
||||
existing_provider_model_names: set[str] = {
|
||||
row[1] for row in existing_associations if row[1]
|
||||
}
|
||||
|
||||
# 获取所有活跃的 GlobalModel(带映射规则)
|
||||
global_models = db.query(GlobalModel).filter(GlobalModel.is_active == True).all()
|
||||
|
||||
allowed_models_set = set(allowed_models)
|
||||
|
||||
for global_model in global_models:
|
||||
# 跳过已关联的
|
||||
if global_model.id in linked_global_model_ids:
|
||||
continue
|
||||
|
||||
# 跳过 provider_model_name 已存在的(避免唯一约束冲突)
|
||||
if global_model.name in existing_provider_model_names:
|
||||
logger.debug(
|
||||
f"Skipping auto-association for GlobalModel {global_model.name}: "
|
||||
f"provider_model_name already exists for Provider {provider.name}"
|
||||
)
|
||||
continue
|
||||
|
||||
# 提取映射规则
|
||||
model_mappings: list[str] = []
|
||||
if global_model.config and isinstance(global_model.config, dict):
|
||||
mappings = global_model.config.get("model_mappings")
|
||||
if isinstance(mappings, list):
|
||||
model_mappings = [m for m in mappings if isinstance(m, str)]
|
||||
|
||||
if not model_mappings:
|
||||
continue
|
||||
|
||||
# 检查是否有映射规则匹配到 Key 白名单
|
||||
matched = False
|
||||
for mapping_pattern in model_mappings:
|
||||
for allowed_model in allowed_models_set:
|
||||
if match_model_with_pattern(mapping_pattern, allowed_model):
|
||||
matched = True
|
||||
break
|
||||
if matched:
|
||||
break
|
||||
|
||||
if not matched:
|
||||
continue
|
||||
|
||||
# 自动创建关联(逐个处理,允许部分成功)
|
||||
try:
|
||||
new_model = Model(
|
||||
provider_id=provider_id,
|
||||
global_model_id=global_model.id,
|
||||
provider_model_name=global_model.name,
|
||||
is_active=True,
|
||||
)
|
||||
db.add(new_model)
|
||||
db.flush()
|
||||
|
||||
# 添加到已存在集合,避免后续循环重复创建
|
||||
existing_provider_model_names.add(global_model.name)
|
||||
|
||||
results["success"].append(
|
||||
{
|
||||
"global_model_id": global_model.id,
|
||||
"global_model_name": global_model.name,
|
||||
"model_id": new_model.id,
|
||||
}
|
||||
)
|
||||
logger.info(
|
||||
f"Auto-associated Provider {provider.name} to GlobalModel {global_model.name} "
|
||||
f"via mapping rule match"
|
||||
)
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
logger.error(
|
||||
f"Failed to auto-associate Provider {provider.name} to GlobalModel {global_model.name}: {e}"
|
||||
)
|
||||
results["errors"].append(
|
||||
{
|
||||
"global_model_id": global_model.id,
|
||||
"global_model_name": global_model.name,
|
||||
"error": str(e),
|
||||
}
|
||||
)
|
||||
|
||||
if results["success"]:
|
||||
db.commit()
|
||||
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def auto_disassociate_provider_by_key_whitelist(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
) -> dict:
|
||||
"""
|
||||
根据 Key 白名单自动解除 Provider 与不再匹配的 GlobalModel 的关联
|
||||
|
||||
当 Key 的 allowed_models 更新后调用此方法,检查所有已关联的 GlobalModel,
|
||||
如果其映射规则不再匹配任何 Key 白名单中的模型,则自动删除关联。
|
||||
|
||||
注意:只删除通过映射规则自动关联的 Model(即 GlobalModel 有 model_mappings 配置的)
|
||||
|
||||
Args:
|
||||
db: 数据库 Session
|
||||
provider_id: Provider ID
|
||||
|
||||
Returns:
|
||||
Dict: 包含 success 和 errors 列表
|
||||
"""
|
||||
from src.core.model_permissions import match_model_with_pattern
|
||||
from src.models.database import Provider, ProviderAPIKey
|
||||
|
||||
results: dict[str, list[dict]] = {
|
||||
"success": [],
|
||||
"errors": [],
|
||||
}
|
||||
|
||||
# 获取 Provider
|
||||
provider = db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
if not provider:
|
||||
logger.warning(f"Provider {provider_id} not found for auto-disassociation")
|
||||
return results
|
||||
|
||||
# 1. 先快速检查是否存在"允许所有模型"的活跃 Key。
|
||||
# 这种情况下无需解除任何关联,避免继续扫描整张 key 表。
|
||||
# 注意:跳过 OAuth Key,OAuth Key 的 allowed_models 由上游动态获取,数量庞大,
|
||||
# 不应参与 disassociate 判定。
|
||||
from src.services.provider_keys.auth_type import OAUTH_AUTH_TYPES
|
||||
|
||||
non_oauth_filter = ProviderAPIKey.auth_type.notin_(OAUTH_AUTH_TYPES)
|
||||
has_unlimited_key = (
|
||||
db.query(ProviderAPIKey.id)
|
||||
.filter(
|
||||
ProviderAPIKey.provider_id == provider_id,
|
||||
ProviderAPIKey.is_active == True,
|
||||
ProviderAPIKey.allowed_models.is_(None),
|
||||
non_oauth_filter,
|
||||
)
|
||||
.limit(1)
|
||||
.first()
|
||||
is not None
|
||||
)
|
||||
if has_unlimited_key:
|
||||
return results
|
||||
|
||||
# 2. 仅查询活跃 Key 的 allowed_models 列,避免把 api_key/auth_config 等大字段整行拉出。
|
||||
allowed_model_rows = (
|
||||
db.query(ProviderAPIKey.allowed_models)
|
||||
.filter(
|
||||
ProviderAPIKey.provider_id == provider_id,
|
||||
ProviderAPIKey.is_active == True,
|
||||
non_oauth_filter,
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
# 如果 Provider 无活跃 Key,不做任何解除(保留现有关联)
|
||||
if not allowed_model_rows:
|
||||
return results
|
||||
|
||||
# 收集所有 Key 的 allowed_models 并集
|
||||
all_allowed_models: set[str] = set()
|
||||
for (allowed_models,) in allowed_model_rows:
|
||||
if isinstance(allowed_models, list) and allowed_models:
|
||||
all_allowed_models.update(m for m in allowed_models if isinstance(m, str))
|
||||
|
||||
# 3. 获取 Provider 当前关联的所有 Model(仅加载判定所需字段)
|
||||
models = (
|
||||
db.query(Model)
|
||||
.options(
|
||||
load_only(Model.id, Model.provider_id, Model.global_model_id),
|
||||
joinedload(Model.global_model).load_only(
|
||||
GlobalModel.id,
|
||||
GlobalModel.name,
|
||||
GlobalModel.config,
|
||||
),
|
||||
)
|
||||
.filter(Model.provider_id == provider_id)
|
||||
.all()
|
||||
)
|
||||
|
||||
# 4. 检查每个 Model 是否还能匹配,收集需要删除的 Model
|
||||
models_to_delete: list[Model] = []
|
||||
|
||||
for model in models:
|
||||
# 跳过 global_model 关系未加载的
|
||||
if not model.global_model:
|
||||
continue
|
||||
|
||||
global_model = cast(GlobalModel, model.global_model)
|
||||
|
||||
# 提取映射规则
|
||||
model_mappings: list[str] = []
|
||||
config = global_model.config
|
||||
if config and isinstance(config, dict):
|
||||
mappings = config.get("model_mappings")
|
||||
if isinstance(mappings, list):
|
||||
model_mappings = [m for m in mappings if isinstance(m, str)]
|
||||
|
||||
# 如果 GlobalModel 没有 model_mappings,跳过(说明不是通过映射自动关联的)
|
||||
if not model_mappings:
|
||||
continue
|
||||
|
||||
# 检查是否有映射规则匹配到任一 allowed_models
|
||||
matched = False
|
||||
for mapping_pattern in model_mappings:
|
||||
for allowed_model in all_allowed_models:
|
||||
if match_model_with_pattern(mapping_pattern, allowed_model):
|
||||
matched = True
|
||||
break
|
||||
if matched:
|
||||
break
|
||||
|
||||
# 如果不再匹配,标记为待删除
|
||||
if not matched:
|
||||
models_to_delete.append(model)
|
||||
|
||||
# 5. 批量删除不再匹配的 Model(全部成功或全部失败)
|
||||
if models_to_delete:
|
||||
try:
|
||||
for model in models_to_delete:
|
||||
global_model = cast(GlobalModel, model.global_model)
|
||||
db.delete(model)
|
||||
results["success"].append(
|
||||
{
|
||||
"model_id": model.id,
|
||||
"global_model_id": global_model.id,
|
||||
"global_model_name": global_model.name,
|
||||
}
|
||||
)
|
||||
logger.info(
|
||||
f"Auto-disassociated Provider {provider.name} from GlobalModel {global_model.name} "
|
||||
f"(no matching allowed_models)"
|
||||
)
|
||||
db.commit()
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
logger.error(f"Failed to auto-disassociate Provider {provider.name}: {e}")
|
||||
# 清空 success,记录整体错误
|
||||
results["success"] = []
|
||||
results["errors"].append(
|
||||
{
|
||||
"provider_id": provider_id,
|
||||
"error": str(e),
|
||||
}
|
||||
)
|
||||
|
||||
return results
|
||||
429
_deprecated_py_src/services/model/mapper.py
Normal file
429
_deprecated_py_src/services/model/mapper.py
Normal file
@@ -0,0 +1,429 @@
|
||||
"""
|
||||
模型映射中间件
|
||||
根据数据库中的配置,将用户请求的模型映射到提供商的实际模型
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy.orm import Session, joinedload
|
||||
|
||||
from src.core.cache_utils import SyncLRUCache
|
||||
from src.core.logger import logger
|
||||
from src.models.claude import ClaudeMessagesRequest
|
||||
from src.models.database import GlobalModel, Model, Provider, ProviderEndpoint
|
||||
from src.services.cache.model_cache import ModelCacheService
|
||||
|
||||
# 模块级共享缓存,所有 ModelMapperMiddleware 实例共用
|
||||
_shared_cache = SyncLRUCache(max_size=1000, ttl=300)
|
||||
|
||||
|
||||
class ModelMapperMiddleware:
|
||||
"""
|
||||
模型映射中间件
|
||||
负责将用户请求的模型名映射到提供商的实际模型名
|
||||
"""
|
||||
|
||||
def __init__(self, db: Session):
|
||||
"""
|
||||
初始化模型映射中间件
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
"""
|
||||
self.db = db
|
||||
|
||||
async def apply_mapping(
|
||||
self, request: ClaudeMessagesRequest, provider: Provider
|
||||
) -> ClaudeMessagesRequest:
|
||||
"""
|
||||
应用模型映射到请求
|
||||
|
||||
Args:
|
||||
request: 原始请求
|
||||
provider: 目标提供商
|
||||
|
||||
Returns:
|
||||
应用映射后的请求
|
||||
"""
|
||||
# 获取请求的模型名
|
||||
source_model = request.model
|
||||
|
||||
# 查找映射
|
||||
mapping = await self.get_mapping(source_model, provider.id)
|
||||
|
||||
if mapping:
|
||||
# 应用映射
|
||||
original_model = request.model
|
||||
request.model = mapping.model.select_provider_model_name()
|
||||
|
||||
logger.debug(
|
||||
f"Applied model mapping for provider {provider.name}: "
|
||||
f"{original_model} -> {request.model}"
|
||||
)
|
||||
else:
|
||||
# 没有找到映射,使用原始模型名
|
||||
logger.debug(
|
||||
f"No model mapping found for {source_model} with provider {provider.name}, "
|
||||
f"forwarding with original model name"
|
||||
)
|
||||
|
||||
return request
|
||||
|
||||
async def get_mapping(self, source_model: str, provider_id: str) -> object | None:
|
||||
"""
|
||||
获取模型映射
|
||||
|
||||
简化后的逻辑:
|
||||
1. 通过 GlobalModel.name 解析 GlobalModel
|
||||
2. 找到 GlobalModel 后,查找该 Provider 的 Model 实现
|
||||
|
||||
Args:
|
||||
source_model: 用户请求的模型名(必须是 GlobalModel.name)
|
||||
provider_id: 提供商ID (UUID)
|
||||
|
||||
Returns:
|
||||
模型映射对象(包含 model 字段),如果没有找到返回None
|
||||
"""
|
||||
# 步骤 1: 规范化模型名称
|
||||
normalized_name = source_model.strip() if isinstance(source_model, str) else ""
|
||||
if not normalized_name:
|
||||
logger.debug("GlobalModel not found: <empty model name>")
|
||||
return None
|
||||
|
||||
# 检查缓存(使用规范化后的名称)
|
||||
cache_key = f"{provider_id}:{normalized_name}"
|
||||
if cache_key in _shared_cache:
|
||||
return _shared_cache[cache_key]
|
||||
|
||||
mapping = None
|
||||
|
||||
global_model = await ModelCacheService.get_global_model_by_name(self.db, normalized_name)
|
||||
|
||||
if not global_model or not global_model.is_active:
|
||||
logger.debug(f"GlobalModel not found or inactive: {normalized_name}")
|
||||
_shared_cache[cache_key] = None
|
||||
return None
|
||||
|
||||
# 步骤 2: 查找该 Provider 是否有实现这个 GlobalModel 的 Model(使用缓存)
|
||||
model = await ModelCacheService.get_model_by_provider_and_global_model(
|
||||
self.db, provider_id, global_model.id
|
||||
)
|
||||
|
||||
if model:
|
||||
# 将 ORM Model 转为无 Session 绑定的实例,避免跨请求缓存导致 DetachedInstanceError
|
||||
from sqlalchemy.orm.session import object_session
|
||||
|
||||
if object_session(model) is not None:
|
||||
model_dict = ModelCacheService._model_to_dict(model)
|
||||
model = ModelCacheService._dict_to_model(model_dict)
|
||||
|
||||
# 创建映射对象
|
||||
mapping = type(
|
||||
"obj",
|
||||
(object,),
|
||||
{
|
||||
"source_model": source_model,
|
||||
"model": model,
|
||||
"is_active": True,
|
||||
"provider_id": provider_id,
|
||||
},
|
||||
)()
|
||||
|
||||
logger.debug(
|
||||
f"Found model mapping: {normalized_name} -> {model.provider_model_name} "
|
||||
f"(provider={provider_id[:8]}...)"
|
||||
)
|
||||
|
||||
# 缓存结果
|
||||
_shared_cache[cache_key] = mapping
|
||||
|
||||
return mapping
|
||||
|
||||
def get_all_mappings(self, provider_id: str) -> list[object]:
|
||||
"""
|
||||
获取提供商的所有可用模型(通过 GlobalModel)
|
||||
|
||||
Args:
|
||||
provider_id: 提供商ID (UUID)
|
||||
|
||||
Returns:
|
||||
模型映射列表
|
||||
"""
|
||||
# 查询该 Provider 的所有活跃 Model(使用 joinedload 避免 N+1)
|
||||
models = (
|
||||
self.db.query(Model)
|
||||
.join(GlobalModel)
|
||||
.options(joinedload(Model.global_model))
|
||||
.filter(
|
||||
Model.provider_id == provider_id,
|
||||
Model.is_active == True,
|
||||
GlobalModel.is_active == True,
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
# 构造兼容的映射对象列表
|
||||
mappings = []
|
||||
for model in models:
|
||||
mapping = type(
|
||||
"obj",
|
||||
(object,),
|
||||
{
|
||||
"source_model": model.global_model.name,
|
||||
"model": model,
|
||||
"is_active": True,
|
||||
"provider_id": provider_id,
|
||||
},
|
||||
)()
|
||||
mappings.append(mapping)
|
||||
|
||||
return mappings
|
||||
|
||||
def get_supported_models(self, provider_id: str) -> list[str]:
|
||||
"""
|
||||
获取提供商支持的所有源模型名
|
||||
|
||||
Args:
|
||||
provider_id: 提供商ID (UUID)
|
||||
|
||||
Returns:
|
||||
支持的模型名列表
|
||||
"""
|
||||
mappings = self.get_all_mappings(provider_id)
|
||||
return [mapping.source_model for mapping in mappings]
|
||||
|
||||
async def validate_request(
|
||||
self, request: ClaudeMessagesRequest, provider: Provider
|
||||
) -> tuple[bool, str | None]:
|
||||
"""
|
||||
验证请求是否符合映射的限制
|
||||
|
||||
Args:
|
||||
request: 请求对象
|
||||
provider: 提供商对象
|
||||
|
||||
Returns:
|
||||
(是否有效, 错误信息)
|
||||
"""
|
||||
mapping = await self.get_mapping(request.model, provider.id)
|
||||
|
||||
if not mapping:
|
||||
# 没有映射,可能是默认支持的模型
|
||||
return True, None
|
||||
|
||||
if not mapping.is_active:
|
||||
return False, f"Model mapping for {request.model} is disabled"
|
||||
|
||||
return True, None
|
||||
|
||||
@staticmethod
|
||||
def clear_cache() -> None:
|
||||
"""清空共享缓存"""
|
||||
_shared_cache.clear()
|
||||
logger.debug("Model mapping cache cleared")
|
||||
|
||||
@staticmethod
|
||||
def refresh_cache(provider_id: str | None = None) -> None:
|
||||
"""
|
||||
刷新缓存
|
||||
|
||||
Args:
|
||||
provider_id: 如果指定,只刷新该提供商的缓存 (UUID)
|
||||
"""
|
||||
if provider_id:
|
||||
keys_to_remove = [
|
||||
key for key in _shared_cache.keys() if key.startswith(f"{provider_id}:")
|
||||
]
|
||||
for key in keys_to_remove:
|
||||
del _shared_cache[key]
|
||||
logger.debug(f"Refreshed cache for provider {provider_id}")
|
||||
else:
|
||||
ModelMapperMiddleware.clear_cache()
|
||||
|
||||
|
||||
class ModelRoutingMiddleware:
|
||||
"""
|
||||
模型路由中间件
|
||||
根据模型名选择合适的提供商
|
||||
"""
|
||||
|
||||
def __init__(self, db: Session):
|
||||
"""
|
||||
初始化模型路由中间件
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
"""
|
||||
self.db = db
|
||||
self.mapper = ModelMapperMiddleware(db)
|
||||
|
||||
def select_provider(
|
||||
self,
|
||||
model_name: str,
|
||||
preferred_provider: str | None = None,
|
||||
allowed_api_formats: list[str] | None = None,
|
||||
request_id: str | None = None,
|
||||
) -> Provider | None:
|
||||
"""
|
||||
根据模型名选择提供商
|
||||
|
||||
Args:
|
||||
model_name: 请求的模型名
|
||||
preferred_provider: 首选提供商名称
|
||||
allowed_api_formats: 允许的API格式列表
|
||||
request_id: 请求ID(用于日志关联)
|
||||
|
||||
Returns:
|
||||
选中的提供商,如果没有找到返回None
|
||||
"""
|
||||
request_prefix = f"ID:{request_id} | " if request_id else ""
|
||||
allowed_norm: set[str] | None = None
|
||||
if allowed_api_formats:
|
||||
from src.services.provider.format import normalize_endpoint_signature
|
||||
|
||||
allowed_norm = {
|
||||
normalize_endpoint_signature(str(fmt))
|
||||
for fmt in allowed_api_formats
|
||||
if isinstance(fmt, str) and fmt
|
||||
}
|
||||
|
||||
# 1. 如果指定了提供商,直接使用
|
||||
if preferred_provider:
|
||||
provider = (
|
||||
self.db.query(Provider)
|
||||
.filter(Provider.name == preferred_provider, Provider.is_active == True)
|
||||
.first()
|
||||
)
|
||||
|
||||
if provider:
|
||||
# 检查API格式 - 从 endpoints 中检查
|
||||
if allowed_norm:
|
||||
has_matching_endpoint = any(
|
||||
ep.is_active
|
||||
and ep.api_format
|
||||
and str(ep.api_format).strip().lower() in allowed_norm
|
||||
for ep in provider.endpoints
|
||||
)
|
||||
if not has_matching_endpoint:
|
||||
logger.warning(
|
||||
f"Specified provider {provider.name} has no active endpoints with allowed API formats ({allowed_api_formats})"
|
||||
)
|
||||
else:
|
||||
logger.debug(
|
||||
f" └─ {request_prefix}使用指定提供商: {provider.name} | 模型:{model_name}"
|
||||
)
|
||||
return provider
|
||||
else:
|
||||
logger.debug(
|
||||
f" └─ {request_prefix}使用指定提供商: {provider.name} | 模型:{model_name}"
|
||||
)
|
||||
return provider
|
||||
else:
|
||||
logger.warning(f"Specified provider {preferred_provider} not found or inactive")
|
||||
|
||||
# 2. 查找优先级最高的活动提供商
|
||||
query = self.db.query(Provider).filter(Provider.is_active == True)
|
||||
|
||||
if allowed_norm:
|
||||
query = (
|
||||
query.join(ProviderEndpoint)
|
||||
.filter(
|
||||
ProviderEndpoint.is_active == True,
|
||||
ProviderEndpoint.api_format.in_(sorted(allowed_norm)),
|
||||
)
|
||||
.distinct()
|
||||
)
|
||||
|
||||
best_provider = query.order_by(Provider.provider_priority.asc(), Provider.id.asc()).first()
|
||||
|
||||
if best_provider:
|
||||
logger.debug(
|
||||
f" └─ {request_prefix}使用优先级最高提供商: {best_provider.name} (priority:{best_provider.provider_priority}) | 模型:{model_name}"
|
||||
)
|
||||
return best_provider
|
||||
|
||||
if allowed_api_formats:
|
||||
logger.error(
|
||||
f"No active providers found with allowed API formats {allowed_api_formats}."
|
||||
)
|
||||
else:
|
||||
logger.error("No active providers found.")
|
||||
return None
|
||||
|
||||
def get_available_models(self) -> dict[str, list[str]]:
|
||||
"""
|
||||
获取所有可用的模型及其提供商
|
||||
|
||||
Returns:
|
||||
字典,键为 GlobalModel.name,值为支持该模型的提供商名列表
|
||||
"""
|
||||
result = {}
|
||||
|
||||
models = (
|
||||
self.db.query(GlobalModel.name, Provider.name)
|
||||
.join(Model, GlobalModel.id == Model.global_model_id)
|
||||
.join(Provider, Model.provider_id == Provider.id)
|
||||
.filter(
|
||||
GlobalModel.is_active == True, Model.is_active == True, Provider.is_active == True
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
for global_model_name, provider_name in models:
|
||||
if global_model_name not in result:
|
||||
result[global_model_name] = []
|
||||
if provider_name not in result[global_model_name]:
|
||||
result[global_model_name].append(provider_name)
|
||||
|
||||
return result
|
||||
|
||||
async def get_cheapest_provider(self, model_name: str) -> Provider | None:
|
||||
"""
|
||||
获取某个模型最便宜的提供商
|
||||
|
||||
Args:
|
||||
model_name: GlobalModel 名称
|
||||
|
||||
Returns:
|
||||
最便宜的提供商
|
||||
"""
|
||||
# 直接查找 GlobalModel
|
||||
global_model = (
|
||||
self.db.query(GlobalModel)
|
||||
.filter(GlobalModel.name == model_name, GlobalModel.is_active == True)
|
||||
.first()
|
||||
)
|
||||
|
||||
if not global_model:
|
||||
return None
|
||||
|
||||
# 查询所有支持该模型的 Provider 及其价格
|
||||
models_with_providers = (
|
||||
self.db.query(Provider, Model)
|
||||
.join(Model, Provider.id == Model.provider_id)
|
||||
.filter(
|
||||
Model.global_model_id == global_model.id,
|
||||
Model.is_active == True,
|
||||
Provider.is_active == True,
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
if not models_with_providers:
|
||||
return None
|
||||
|
||||
# 按总价格排序
|
||||
cheapest = min(
|
||||
models_with_providers,
|
||||
key=lambda x: x[1].get_effective_input_price() + x[1].get_effective_output_price(),
|
||||
)
|
||||
|
||||
provider = cheapest[0]
|
||||
model = cheapest[1]
|
||||
|
||||
logger.debug(
|
||||
f"Selected cheapest provider {provider.name} for model {model_name} "
|
||||
f"(input: ${model.get_effective_input_price()}/M, output: ${model.get_effective_output_price()}/M)"
|
||||
)
|
||||
|
||||
return provider
|
||||
50
_deprecated_py_src/services/model/pricing_strategy.py
Normal file
50
_deprecated_py_src/services/model/pricing_strategy.py
Normal file
@@ -0,0 +1,50 @@
|
||||
"""
|
||||
计费相关数据类
|
||||
|
||||
定义计费计算所需的数据结构。
|
||||
实际的计费能力已收敛到 core.api_format 注册表,避免依赖 API Adapter。
|
||||
|
||||
数据类:
|
||||
- UsageTokens: 请求的 token 使用量
|
||||
- PricingConfig: 价格配置
|
||||
- CostResult: 计费结果
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class UsageTokens:
|
||||
"""请求的 token 使用量"""
|
||||
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cache_creation_input_tokens: int = 0
|
||||
cache_read_input_tokens: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class PricingConfig:
|
||||
"""价格配置"""
|
||||
|
||||
input_price_per_1m: float = 0.0
|
||||
output_price_per_1m: float = 0.0
|
||||
cache_creation_price_per_1m: float | None = None
|
||||
cache_read_price_per_1m: float | None = None
|
||||
price_per_request: float | None = None
|
||||
tiered_pricing: dict | None = None
|
||||
cache_ttl_minutes: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class CostResult:
|
||||
"""计费结果"""
|
||||
|
||||
input_cost: float = 0.0
|
||||
output_cost: float = 0.0
|
||||
cache_creation_cost: float = 0.0
|
||||
cache_read_cost: float = 0.0
|
||||
cache_cost: float = 0.0
|
||||
request_cost: float = 0.0
|
||||
total_cost: float = 0.0
|
||||
tier_index: int | None = None # 命中的阶梯索引
|
||||
476
_deprecated_py_src/services/model/service.py
Normal file
476
_deprecated_py_src/services/model/service.py
Normal file
@@ -0,0 +1,476 @@
|
||||
"""
|
||||
模型管理服务
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy import and_, func
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.exceptions import InvalidRequestException, NotFoundException
|
||||
from src.core.logger import logger
|
||||
from src.models.api import ModelCreate, ModelResponse, ModelUpdate
|
||||
from src.models.database import Model, Provider
|
||||
from src.services.cache.invalidation import get_cache_invalidation_service
|
||||
from src.services.cache.model_cache import ModelCacheService
|
||||
from src.services.cache.model_list_cache import invalidate_models_list_cache
|
||||
from src.utils.async_utils import safe_create_task
|
||||
|
||||
|
||||
class ModelService:
|
||||
"""模型管理服务"""
|
||||
|
||||
@staticmethod
|
||||
def create_model(db: Session, provider_id: str, model_data: ModelCreate) -> Model:
|
||||
"""创建模型"""
|
||||
# 检查提供商是否存在
|
||||
provider = db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
if not provider:
|
||||
raise NotFoundException(f"提供商 {provider_id} 不存在")
|
||||
|
||||
# 检查同一提供商下是否已存在同名模型
|
||||
existing = (
|
||||
db.query(Model)
|
||||
.filter(
|
||||
and_(
|
||||
Model.provider_id == provider_id,
|
||||
Model.provider_model_name == model_data.provider_model_name,
|
||||
)
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if existing:
|
||||
raise InvalidRequestException(
|
||||
f"提供商 {provider.name} 下已存在模型 {model_data.provider_model_name}"
|
||||
)
|
||||
|
||||
try:
|
||||
model = Model(
|
||||
provider_id=provider_id,
|
||||
global_model_id=model_data.global_model_id,
|
||||
provider_model_name=model_data.provider_model_name,
|
||||
provider_model_mappings=model_data.provider_model_mappings,
|
||||
price_per_request=model_data.price_per_request,
|
||||
tiered_pricing=model_data.tiered_pricing,
|
||||
supports_vision=model_data.supports_vision,
|
||||
supports_function_calling=model_data.supports_function_calling,
|
||||
supports_streaming=model_data.supports_streaming,
|
||||
supports_extended_thinking=model_data.supports_extended_thinking,
|
||||
is_active=model_data.is_active if model_data.is_active is not None else True,
|
||||
config=model_data.config,
|
||||
)
|
||||
db.add(model)
|
||||
db.commit()
|
||||
db.refresh(model)
|
||||
# 显式加载 global_model 关系
|
||||
from sqlalchemy.orm import joinedload
|
||||
|
||||
model = (
|
||||
db.query(Model)
|
||||
.options(joinedload(Model.global_model))
|
||||
.filter(Model.id == model.id)
|
||||
.first()
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"创建模型成功: provider={provider.name}, model={model.provider_model_name}, global_model_id={model.global_model_id}"
|
||||
)
|
||||
|
||||
# 清除 Redis 缓存(异步执行,不阻塞返回)
|
||||
# 重要:新增模型可能需要清除 resolver 的 NOT_FOUND 负缓存(global_model:resolve:*),
|
||||
# 否则请求链路在 TTL 内可能无法立刻解析到新模型。
|
||||
safe_create_task(
|
||||
ModelCacheService.invalidate_model_cache(
|
||||
model_id=model.id,
|
||||
provider_id=model.provider_id,
|
||||
global_model_id=model.global_model_id,
|
||||
provider_model_name=model.provider_model_name,
|
||||
provider_model_mappings=model.provider_model_mappings,
|
||||
)
|
||||
)
|
||||
|
||||
# 清除内存缓存(ModelMapperMiddleware 实例)
|
||||
if model.provider_id and model.global_model_id:
|
||||
cache_service = get_cache_invalidation_service()
|
||||
cache_service.on_model_changed(model.provider_id, model.global_model_id)
|
||||
|
||||
# 清除 /v1/models 列表缓存
|
||||
safe_create_task(invalidate_models_list_cache())
|
||||
|
||||
return model
|
||||
|
||||
except IntegrityError as e:
|
||||
db.rollback()
|
||||
logger.error(f"创建模型失败: {str(e)}")
|
||||
raise InvalidRequestException("创建模型失败,请检查输入数据")
|
||||
|
||||
@staticmethod
|
||||
def get_model(db: Session, model_id: str) -> Model: # UUID
|
||||
"""获取模型详情"""
|
||||
from sqlalchemy.orm import joinedload
|
||||
|
||||
model = (
|
||||
db.query(Model)
|
||||
.options(joinedload(Model.global_model))
|
||||
.filter(Model.id == model_id)
|
||||
.first()
|
||||
)
|
||||
if not model:
|
||||
raise NotFoundException(f"模型 {model_id} 不存在")
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
def get_models_by_provider(
|
||||
db: Session,
|
||||
provider_id: str, # UUID
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
is_active: bool | None = None,
|
||||
) -> list[Model]:
|
||||
"""获取提供商的模型列表"""
|
||||
from sqlalchemy.orm import joinedload
|
||||
|
||||
query = (
|
||||
db.query(Model)
|
||||
.options(joinedload(Model.global_model))
|
||||
.filter(Model.provider_id == provider_id)
|
||||
)
|
||||
|
||||
if is_active is not None:
|
||||
query = query.filter(Model.is_active == is_active)
|
||||
|
||||
# 按创建时间排序
|
||||
query = query.order_by(Model.created_at.desc())
|
||||
|
||||
return query.offset(skip).limit(limit).all()
|
||||
|
||||
@staticmethod
|
||||
def get_all_models(
|
||||
db: Session,
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
is_active: bool | None = None,
|
||||
category: str | None = None,
|
||||
) -> list[Model]:
|
||||
"""获取所有模型列表"""
|
||||
query = db.query(Model)
|
||||
|
||||
if is_active is not None:
|
||||
query = query.filter(Model.is_active == is_active)
|
||||
|
||||
# 按提供商和创建时间排序
|
||||
query = query.order_by(Model.provider_id, Model.created_at.desc())
|
||||
|
||||
return query.offset(skip).limit(limit).all()
|
||||
|
||||
@staticmethod
|
||||
def update_model(db: Session, model_id: str, model_data: ModelUpdate) -> Model: # UUID
|
||||
"""更新模型"""
|
||||
model = db.query(Model).filter(Model.id == model_id).first()
|
||||
if not model:
|
||||
raise NotFoundException(f"模型 {model_id} 不存在")
|
||||
|
||||
# 保存旧的映射,用于清除缓存
|
||||
old_global_model_id = model.global_model_id
|
||||
old_provider_model_name = model.provider_model_name
|
||||
old_provider_model_mappings = model.provider_model_mappings
|
||||
|
||||
# 更新字段
|
||||
update_data = model_data.model_dump(exclude_unset=True)
|
||||
|
||||
# 添加调试日志
|
||||
logger.debug(f"更新模型 {model_id} 收到的数据: {update_data}")
|
||||
logger.debug(
|
||||
f"更新前的 supports_vision: {model.supports_vision}, supports_function_calling: {model.supports_function_calling}, supports_extended_thinking: {model.supports_extended_thinking}"
|
||||
)
|
||||
|
||||
for field, value in update_data.items():
|
||||
setattr(model, field, value)
|
||||
|
||||
logger.debug(
|
||||
f"更新后的 supports_vision: {model.supports_vision}, supports_function_calling: {model.supports_function_calling}, supports_extended_thinking: {model.supports_extended_thinking}"
|
||||
)
|
||||
|
||||
try:
|
||||
db.commit()
|
||||
db.refresh(model)
|
||||
|
||||
# 清除 Redis 缓存(异步执行,不阻塞返回)
|
||||
# 先清除旧的映射缓存
|
||||
safe_create_task(
|
||||
ModelCacheService.invalidate_model_cache(
|
||||
model_id=model.id,
|
||||
provider_id=model.provider_id,
|
||||
global_model_id=old_global_model_id,
|
||||
provider_model_name=old_provider_model_name,
|
||||
provider_model_mappings=old_provider_model_mappings,
|
||||
)
|
||||
)
|
||||
# 再清除新的映射缓存(如果有变化,包括 global_model_id 变更)
|
||||
if (
|
||||
model.provider_model_name != old_provider_model_name
|
||||
or model.provider_model_mappings != old_provider_model_mappings
|
||||
or model.global_model_id != old_global_model_id
|
||||
):
|
||||
safe_create_task(
|
||||
ModelCacheService.invalidate_model_cache(
|
||||
model_id=model.id,
|
||||
provider_id=model.provider_id,
|
||||
global_model_id=model.global_model_id,
|
||||
provider_model_name=model.provider_model_name,
|
||||
provider_model_mappings=model.provider_model_mappings,
|
||||
)
|
||||
)
|
||||
|
||||
# 清除内存缓存(ModelMapperMiddleware 实例)
|
||||
if model.provider_id:
|
||||
cache_service = get_cache_invalidation_service()
|
||||
cache_service.on_model_changed(model.provider_id, model.global_model_id)
|
||||
|
||||
# 清除 /v1/models 列表缓存
|
||||
safe_create_task(invalidate_models_list_cache())
|
||||
|
||||
logger.info(
|
||||
f"更新模型成功: id={model_id}, 最终 supports_vision: {model.supports_vision}, supports_function_calling: {model.supports_function_calling}, supports_extended_thinking: {model.supports_extended_thinking}"
|
||||
)
|
||||
return model
|
||||
except IntegrityError as e:
|
||||
db.rollback()
|
||||
logger.error(f"更新模型失败: {str(e)}")
|
||||
raise InvalidRequestException("更新模型失败,请检查输入数据")
|
||||
|
||||
@staticmethod
|
||||
def delete_model(db: Session, model_id: str) -> None: # UUID
|
||||
"""删除模型
|
||||
|
||||
删除逻辑:
|
||||
- Model 只是 Provider 对 GlobalModel 的实现,删除不影响 GlobalModel
|
||||
- 检查是否是该 GlobalModel 的最后一个实现(如果是,警告但允许删除)
|
||||
"""
|
||||
model = db.query(Model).filter(Model.id == model_id).first()
|
||||
if not model:
|
||||
raise NotFoundException(f"模型 {model_id} 不存在")
|
||||
|
||||
# 检查这是否是该 GlobalModel 的最后一个关联提供商
|
||||
if model.global_model_id:
|
||||
other_implementations = int(
|
||||
db.query(func.count(Model.id))
|
||||
.filter(
|
||||
Model.global_model_id == model.global_model_id,
|
||||
Model.id != model_id,
|
||||
Model.is_active == True,
|
||||
)
|
||||
.scalar()
|
||||
or 0
|
||||
)
|
||||
|
||||
if other_implementations == 0:
|
||||
logger.warning(
|
||||
f"警告:删除模型 {model_id}(Provider: {model.provider_id[:8]}...)后,"
|
||||
f"GlobalModel '{model.global_model_id}' 将没有任何活跃的关联提供商"
|
||||
)
|
||||
|
||||
# 保存缓存清除所需的信息(删除后无法访问)
|
||||
cache_info = {
|
||||
"model_id": model.id,
|
||||
"provider_id": model.provider_id,
|
||||
"global_model_id": model.global_model_id,
|
||||
"provider_model_name": model.provider_model_name,
|
||||
"provider_model_mappings": model.provider_model_mappings,
|
||||
}
|
||||
|
||||
try:
|
||||
db.delete(model)
|
||||
db.commit()
|
||||
|
||||
# 清除 Redis 缓存
|
||||
safe_create_task(
|
||||
ModelCacheService.invalidate_model_cache(
|
||||
model_id=cache_info["model_id"],
|
||||
provider_id=cache_info["provider_id"],
|
||||
global_model_id=cache_info["global_model_id"],
|
||||
provider_model_name=cache_info["provider_model_name"],
|
||||
provider_model_mappings=cache_info["provider_model_mappings"],
|
||||
)
|
||||
)
|
||||
|
||||
# 清除内存缓存
|
||||
if cache_info["provider_id"]:
|
||||
cache_service = get_cache_invalidation_service()
|
||||
cache_service.on_model_changed(
|
||||
cache_info["provider_id"], cache_info["global_model_id"]
|
||||
)
|
||||
|
||||
# 清除 /v1/models 列表缓存
|
||||
safe_create_task(invalidate_models_list_cache())
|
||||
|
||||
logger.info(
|
||||
f"删除模型成功: id={model_id}, provider_model_name={cache_info['provider_model_name']}, "
|
||||
f"global_model_id={cache_info['global_model_id'][:8] if cache_info['global_model_id'] else 'None'}..."
|
||||
)
|
||||
except Exception as e:
|
||||
db.rollback()
|
||||
logger.error(f"删除模型失败: {str(e)}")
|
||||
raise InvalidRequestException("删除模型失败")
|
||||
|
||||
@staticmethod
|
||||
def toggle_model_availability(db: Session, model_id: str, is_available: bool) -> Model: # UUID
|
||||
"""切换模型可用状态"""
|
||||
model = db.query(Model).filter(Model.id == model_id).first()
|
||||
if not model:
|
||||
raise NotFoundException(f"模型 {model_id} 不存在")
|
||||
|
||||
model.is_available = is_available
|
||||
db.commit()
|
||||
db.refresh(model)
|
||||
|
||||
# 清除 Redis 缓存
|
||||
safe_create_task(
|
||||
ModelCacheService.invalidate_model_cache(
|
||||
model_id=model.id,
|
||||
provider_id=model.provider_id,
|
||||
global_model_id=model.global_model_id,
|
||||
provider_model_name=model.provider_model_name,
|
||||
provider_model_mappings=model.provider_model_mappings,
|
||||
)
|
||||
)
|
||||
|
||||
# 清除内存缓存(ModelMapperMiddleware 实例)
|
||||
if model.provider_id:
|
||||
cache_service = get_cache_invalidation_service()
|
||||
cache_service.on_model_changed(model.provider_id, model.global_model_id)
|
||||
|
||||
# 清除 /v1/models 列表缓存
|
||||
safe_create_task(invalidate_models_list_cache())
|
||||
|
||||
status = "可用" if is_available else "不可用"
|
||||
logger.info(f"更新模型可用状态: id={model_id}, status={status}")
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
def get_model_by_name(db: Session, provider_id: str, model_name: str) -> Model | None:
|
||||
"""根据 provider_model_name 获取模型"""
|
||||
return (
|
||||
db.query(Model)
|
||||
.filter(and_(Model.provider_id == provider_id, Model.provider_model_name == model_name))
|
||||
.first()
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def batch_create_models(
|
||||
db: Session, provider_id: str, models_data: list[ModelCreate]
|
||||
) -> list[Model]: # UUID
|
||||
"""批量创建模型"""
|
||||
# 检查提供商是否存在
|
||||
provider = db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
if not provider:
|
||||
raise NotFoundException(f"提供商 {provider_id} 不存在")
|
||||
|
||||
created_models = []
|
||||
for model_data in models_data:
|
||||
# 检查是否已存在
|
||||
existing = (
|
||||
db.query(Model)
|
||||
.filter(
|
||||
and_(
|
||||
Model.provider_id == provider_id,
|
||||
Model.provider_model_name == model_data.provider_model_name,
|
||||
)
|
||||
)
|
||||
.first()
|
||||
)
|
||||
|
||||
if existing:
|
||||
logger.warning(f"模型 {model_data.provider_model_name} 已存在,跳过创建")
|
||||
continue
|
||||
|
||||
model = Model(
|
||||
provider_id=provider_id,
|
||||
global_model_id=model_data.global_model_id,
|
||||
provider_model_name=model_data.provider_model_name,
|
||||
provider_model_mappings=model_data.provider_model_mappings,
|
||||
price_per_request=model_data.price_per_request,
|
||||
tiered_pricing=model_data.tiered_pricing,
|
||||
supports_vision=model_data.supports_vision,
|
||||
supports_function_calling=model_data.supports_function_calling,
|
||||
supports_streaming=model_data.supports_streaming,
|
||||
supports_extended_thinking=model_data.supports_extended_thinking,
|
||||
is_active=model_data.is_active,
|
||||
config=model_data.config,
|
||||
)
|
||||
db.add(model)
|
||||
created_models.append(model)
|
||||
|
||||
if created_models:
|
||||
try:
|
||||
db.commit()
|
||||
for model in created_models:
|
||||
db.refresh(model)
|
||||
logger.info(f"批量创建 {len(created_models)} 个模型成功")
|
||||
|
||||
# 清除 Redis 缓存(异步执行,不阻塞返回)
|
||||
# 逐个清除 resolver 的映射缓存,避免 NOT_FOUND 负缓存阻塞新模型生效。
|
||||
for model in created_models:
|
||||
safe_create_task(
|
||||
ModelCacheService.invalidate_model_cache(
|
||||
model_id=model.id,
|
||||
provider_id=model.provider_id,
|
||||
global_model_id=model.global_model_id,
|
||||
provider_model_name=model.provider_model_name,
|
||||
provider_model_mappings=model.provider_model_mappings,
|
||||
)
|
||||
)
|
||||
|
||||
# 清除内存缓存(ModelMapperMiddleware 实例)
|
||||
cache_service = get_cache_invalidation_service()
|
||||
cache_service.on_model_changed(provider_id, created_models[0].global_model_id)
|
||||
|
||||
# 清除 /v1/models 列表缓存
|
||||
safe_create_task(invalidate_models_list_cache())
|
||||
except IntegrityError as e:
|
||||
db.rollback()
|
||||
logger.error(f"批量创建模型失败: {str(e)}")
|
||||
raise InvalidRequestException("批量创建模型失败")
|
||||
|
||||
return created_models
|
||||
|
||||
@staticmethod
|
||||
def convert_to_response(model: Model) -> ModelResponse:
|
||||
"""转换为响应模型(从 GlobalModel 获取显示信息和默认值)"""
|
||||
return ModelResponse(
|
||||
id=model.id,
|
||||
provider_id=model.provider_id,
|
||||
global_model_id=model.global_model_id,
|
||||
provider_model_name=model.provider_model_name,
|
||||
provider_model_mappings=model.provider_model_mappings,
|
||||
# 原始配置值(可能为空)
|
||||
price_per_request=model.price_per_request,
|
||||
tiered_pricing=model.tiered_pricing,
|
||||
supports_vision=model.supports_vision,
|
||||
supports_function_calling=model.supports_function_calling,
|
||||
supports_streaming=model.supports_streaming,
|
||||
supports_extended_thinking=model.supports_extended_thinking,
|
||||
supports_image_generation=model.supports_image_generation,
|
||||
# 有效值(合并 Model 和 GlobalModel 默认值)
|
||||
effective_tiered_pricing=model.get_effective_tiered_pricing(),
|
||||
effective_input_price=model.get_effective_input_price(),
|
||||
effective_output_price=model.get_effective_output_price(),
|
||||
effective_price_per_request=model.get_effective_price_per_request(),
|
||||
effective_supports_vision=model.get_effective_supports_vision(),
|
||||
effective_supports_function_calling=model.get_effective_supports_function_calling(),
|
||||
effective_supports_streaming=model.get_effective_supports_streaming(),
|
||||
effective_supports_extended_thinking=model.get_effective_supports_extended_thinking(),
|
||||
effective_supports_image_generation=model.get_effective_supports_image_generation(),
|
||||
is_active=model.is_active,
|
||||
is_available=model.is_available if model.is_available is not None else True,
|
||||
created_at=model.created_at,
|
||||
updated_at=model.updated_at,
|
||||
# GlobalModel 信息(如果存在)
|
||||
global_model_name=model.global_model.name if model.global_model else None,
|
||||
global_model_display_name=(
|
||||
model.global_model.display_name if model.global_model else None
|
||||
),
|
||||
# 有效配置(合并 Model 和 GlobalModel 的 config)
|
||||
effective_config=model.get_effective_config(),
|
||||
)
|
||||
515
_deprecated_py_src/services/model/upstream_fetcher.py
Normal file
515
_deprecated_py_src/services/model/upstream_fetcher.py
Normal file
@@ -0,0 +1,515 @@
|
||||
"""
|
||||
上游模型获取公共模块
|
||||
|
||||
提供从上游 API 获取模型列表的公共函数;通用 api_format 抓取能力统一来自 core.api_format.capabilities,供以下场景使用:
|
||||
- 定时任务自动获取(fetch_scheduler.py)
|
||||
- 管理后台手动查询(provider_query.py)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import Iterable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Awaitable, Callable
|
||||
|
||||
import httpx
|
||||
|
||||
from src.config.settings import config
|
||||
from src.core.api_format import get_extra_headers_from_endpoint
|
||||
from src.core.api_format.capabilities import BROWSER_FINGERPRINT_HEADERS
|
||||
from src.core.api_format.headers import build_adapter_headers_for_endpoint
|
||||
from src.core.logger import logger
|
||||
|
||||
# 并发请求限制
|
||||
MAX_CONCURRENT_REQUESTS = 5
|
||||
|
||||
# 模型获取格式优先级:同族内优先使用 chat 端点,若无则回退到 cli 端点
|
||||
MODEL_FETCH_FORMAT_PRIORITY: list[tuple[str, ...]] = [
|
||||
("openai:chat", "openai:cli", "openai:compact"),
|
||||
("claude:chat", "claude:cli"),
|
||||
("gemini:chat", "gemini:cli"),
|
||||
]
|
||||
|
||||
# Return tuple signature:
|
||||
# (models, errors, has_success, upstream_metadata)
|
||||
_ModelsFetcher = Callable[
|
||||
["UpstreamModelsFetchContext", float],
|
||||
Awaitable[tuple[list[dict], list[str], bool, dict[str, Any] | None]],
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class EndpointFetchConfig:
|
||||
"""端点获取配置(纯数据,不依赖 DB session)。
|
||||
|
||||
从 ProviderEndpoint ORM 对象提取必要字段,确保在 DB session 关闭后
|
||||
仍可安全使用(避免 DetachedInstanceError)。
|
||||
"""
|
||||
|
||||
base_url: str
|
||||
extra_headers: dict[str, str] | None = None
|
||||
|
||||
|
||||
def build_format_to_config(endpoints: Iterable[Any]) -> dict[str, EndpointFetchConfig]:
|
||||
"""将活跃的 ProviderEndpoint 转换为 api_format -> EndpointFetchConfig 映射。
|
||||
|
||||
应在 DB session 活跃时调用,提取 ORM 对象上的 base_url 和 header_rules,
|
||||
转换为 session 无关的纯数据结构。
|
||||
"""
|
||||
result: dict[str, EndpointFetchConfig] = {}
|
||||
for ep in endpoints:
|
||||
if not getattr(ep, "is_active", False):
|
||||
continue
|
||||
result[ep.api_format] = EndpointFetchConfig(
|
||||
base_url=ep.base_url,
|
||||
extra_headers=get_extra_headers_from_endpoint(ep),
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class UpstreamModelsFetchContext:
|
||||
"""上游模型获取上下文(Key 级别)。"""
|
||||
|
||||
provider_type: str
|
||||
api_key_value: str
|
||||
format_to_endpoint: dict[str, EndpointFetchConfig]
|
||||
proxy_config: dict[str, Any] | None = None
|
||||
auth_config: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class UpstreamModelsFetcherRegistry:
|
||||
"""按 provider_type 注册上游模型获取策略,避免到处写特判。"""
|
||||
|
||||
_fetchers: dict[str, _ModelsFetcher] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, *, provider_types: list[str], fetcher: _ModelsFetcher) -> None:
|
||||
for pt in provider_types:
|
||||
if not pt:
|
||||
continue
|
||||
cls._fetchers[pt.lower()] = fetcher
|
||||
|
||||
@classmethod
|
||||
def get(cls, provider_type: str) -> _ModelsFetcher | None:
|
||||
if not provider_type:
|
||||
return None
|
||||
return cls._fetchers.get(provider_type.lower())
|
||||
|
||||
|
||||
async def _fetch_models_default(
|
||||
ctx: UpstreamModelsFetchContext,
|
||||
timeout_seconds: float,
|
||||
) -> tuple[list[dict], list[str], bool, dict[str, Any] | None]:
|
||||
endpoint_configs = build_all_format_configs(ctx.api_key_value, ctx.format_to_endpoint)
|
||||
models, errors, has_success = await fetch_models_from_endpoints(
|
||||
endpoint_configs, timeout=timeout_seconds, proxy_config=ctx.proxy_config
|
||||
)
|
||||
return models, errors, has_success, None
|
||||
|
||||
|
||||
async def fetch_models_for_key(
|
||||
ctx: UpstreamModelsFetchContext,
|
||||
*,
|
||||
timeout_seconds: float = 30.0,
|
||||
) -> tuple[list[dict], list[str], bool, dict[str, Any] | None]:
|
||||
"""统一入口:按 provider_type 选择策略获取模型列表(可附带 upstream_metadata)。"""
|
||||
# Ensure provider plugins (including custom model fetchers) are registered.
|
||||
from src.services.provider.envelope import ensure_providers_bootstrapped
|
||||
|
||||
ensure_providers_bootstrapped(provider_types=[ctx.provider_type] if ctx.provider_type else None)
|
||||
|
||||
fetcher = UpstreamModelsFetcherRegistry.get(ctx.provider_type) or _fetch_models_default
|
||||
return await fetcher(ctx, timeout_seconds)
|
||||
|
||||
|
||||
def merge_upstream_metadata(
|
||||
current: dict[str, Any] | None,
|
||||
incoming: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""合并上游元数据,对 quota_by_model 做模型级深度合并。
|
||||
|
||||
以上游返回为准(全量替换),不再保留上游已下架的模型。
|
||||
仅对上游返回的模型补充旧数据中的 reset_time(当新数据缺少时)。
|
||||
"""
|
||||
merged: dict[str, Any] = dict(current) if isinstance(current, dict) else {}
|
||||
for ns_key, ns_val in incoming.items():
|
||||
old_ns = merged.get(ns_key)
|
||||
if (
|
||||
isinstance(ns_val, dict)
|
||||
and isinstance(old_ns, dict)
|
||||
and "quota_by_model" in ns_val
|
||||
and "quota_by_model" in old_ns
|
||||
):
|
||||
old_qbm = old_ns["quota_by_model"]
|
||||
new_qbm = ns_val["quota_by_model"]
|
||||
if isinstance(old_qbm, dict) and isinstance(new_qbm, dict):
|
||||
# 保留新数据中已有模型的旧 reset_time
|
||||
for model_id, new_info in new_qbm.items():
|
||||
if not isinstance(new_info, dict):
|
||||
continue
|
||||
old_info = old_qbm.get(model_id)
|
||||
if (
|
||||
isinstance(old_info, dict)
|
||||
and "reset_time" in old_info
|
||||
and "reset_time" not in new_info
|
||||
):
|
||||
new_info["reset_time"] = old_info["reset_time"]
|
||||
merged[ns_key] = ns_val
|
||||
return merged
|
||||
|
||||
|
||||
def build_all_format_configs(
|
||||
api_key_value: str,
|
||||
format_to_endpoint: dict[str, EndpointFetchConfig],
|
||||
) -> list[dict]:
|
||||
"""
|
||||
构建所有 API 格式的端点配置
|
||||
|
||||
只对实际配置了端点的格式构建请求配置,不同端点的 base_url 可能不同,
|
||||
不应使用某个端点的 base_url 去尝试其他格式。
|
||||
|
||||
Args:
|
||||
api_key_value: 解密后的 API Key
|
||||
format_to_endpoint: API 格式到 EndpointFetchConfig 的映射
|
||||
|
||||
Returns:
|
||||
端点配置列表,每个配置包含 api_key, base_url, api_format, extra_headers
|
||||
"""
|
||||
if not format_to_endpoint:
|
||||
return []
|
||||
|
||||
# 同族内优先使用 chat 端点,若无则回退到 cli 端点
|
||||
configs: list[dict] = []
|
||||
for candidates in MODEL_FETCH_FORMAT_PRIORITY:
|
||||
fmt = next((f for f in candidates if f in format_to_endpoint), None)
|
||||
if fmt is not None:
|
||||
cfg = format_to_endpoint[fmt]
|
||||
|
||||
base_url = str(getattr(cfg, "base_url", "") or "")
|
||||
extra_headers: dict[str, str] | None
|
||||
|
||||
if isinstance(cfg, EndpointFetchConfig):
|
||||
extra_headers = cfg.extra_headers
|
||||
else:
|
||||
# 允许直接传递类似 ProviderEndpoint 的对象(测试/独立使用场景)。
|
||||
candidate_extra = getattr(cfg, "extra_headers", None)
|
||||
if isinstance(candidate_extra, dict):
|
||||
extra_headers = {str(k): str(v) for k, v in candidate_extra.items() if k}
|
||||
else:
|
||||
extra_headers = get_extra_headers_from_endpoint(cfg)
|
||||
|
||||
configs.append(
|
||||
{
|
||||
"api_key": api_key_value,
|
||||
"base_url": base_url,
|
||||
"api_format": fmt,
|
||||
"extra_headers": extra_headers,
|
||||
}
|
||||
)
|
||||
return configs
|
||||
|
||||
|
||||
async def _build_models_proxy_snapshot(
|
||||
proxy_config: dict[str, Any] | None,
|
||||
) -> Any:
|
||||
from src.services.request.execution_runtime_plan import build_proxy_snapshot
|
||||
|
||||
return await build_proxy_snapshot(proxy_config, label="upstream model fetch")
|
||||
|
||||
|
||||
def _build_model_fetch_url(
|
||||
api_format: str,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
after_id: str | None = None,
|
||||
) -> str:
|
||||
normalized_api_format = str(api_format or "").strip().lower()
|
||||
base_url = str(base_url or "").rstrip("/")
|
||||
|
||||
if normalized_api_format.startswith("gemini:"):
|
||||
if base_url.endswith("/v1beta"):
|
||||
return f"{base_url}/models?key={api_key}"
|
||||
return f"{base_url}/v1beta/models?key={api_key}"
|
||||
|
||||
if base_url.endswith("/v1"):
|
||||
url = f"{base_url}/models"
|
||||
else:
|
||||
url = f"{base_url}/v1/models"
|
||||
|
||||
if normalized_api_format.startswith("claude:"):
|
||||
if after_id:
|
||||
return f"{url}?limit=100&after_id={after_id}"
|
||||
return f"{url}?limit=100"
|
||||
|
||||
return url
|
||||
|
||||
|
||||
def _build_model_fetch_headers(
|
||||
api_format: str,
|
||||
api_key: str,
|
||||
extra_headers: dict[str, str] | None,
|
||||
) -> dict[str, str]:
|
||||
normalized_api_format = str(api_format or "").strip().lower()
|
||||
|
||||
if normalized_api_format == "openai:cli":
|
||||
headers = {"User-Agent": config.internal_user_agent_openai_cli}
|
||||
if extra_headers:
|
||||
headers.update(extra_headers)
|
||||
return build_adapter_headers_for_endpoint("openai:chat", api_key, headers)
|
||||
|
||||
if normalized_api_format == "openai:compact":
|
||||
headers = {"User-Agent": config.internal_user_agent_openai_cli}
|
||||
if extra_headers:
|
||||
headers.update(extra_headers)
|
||||
return build_adapter_headers_for_endpoint("openai:chat", api_key, headers)
|
||||
|
||||
if normalized_api_format == "claude:cli":
|
||||
headers = {"User-Agent": config.internal_user_agent_claude_cli}
|
||||
if extra_headers:
|
||||
headers.update(extra_headers)
|
||||
return build_adapter_headers_for_endpoint(normalized_api_format, api_key, headers)
|
||||
|
||||
if normalized_api_format == "claude:chat":
|
||||
headers = build_adapter_headers_for_endpoint(normalized_api_format, api_key, extra_headers)
|
||||
if "authorization" not in {str(k).lower() for k in headers}:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
return headers
|
||||
|
||||
if normalized_api_format == "gemini:cli":
|
||||
headers = {
|
||||
**BROWSER_FINGERPRINT_HEADERS,
|
||||
"User-Agent": config.internal_user_agent_gemini_cli,
|
||||
}
|
||||
if extra_headers:
|
||||
headers.update(extra_headers)
|
||||
return headers
|
||||
|
||||
if normalized_api_format == "gemini:chat":
|
||||
headers = {**BROWSER_FINGERPRINT_HEADERS}
|
||||
if extra_headers:
|
||||
headers.update(extra_headers)
|
||||
return headers
|
||||
|
||||
return build_adapter_headers_for_endpoint(normalized_api_format, api_key, extra_headers)
|
||||
|
||||
|
||||
def _parse_model_fetch_response(
|
||||
api_format: str,
|
||||
payload: Any,
|
||||
) -> tuple[list[dict[str, Any]], bool, str | None]:
|
||||
normalized_api_format = str(api_format or "").strip().lower()
|
||||
|
||||
if normalized_api_format.startswith("gemini:"):
|
||||
if isinstance(payload, dict) and isinstance(payload.get("models"), list):
|
||||
out: list[dict[str, Any]] = []
|
||||
for model in payload["models"]:
|
||||
if not isinstance(model, dict):
|
||||
continue
|
||||
out.append(
|
||||
{
|
||||
"id": str(model.get("name", "")).replace("models/", ""),
|
||||
"owned_by": "google",
|
||||
"display_name": model.get("displayName", ""),
|
||||
"api_format": normalized_api_format,
|
||||
}
|
||||
)
|
||||
return out, False, None
|
||||
return [], False, None
|
||||
|
||||
page_models: list[dict[str, Any]] = []
|
||||
has_more = False
|
||||
if isinstance(payload, dict) and isinstance(payload.get("data"), list):
|
||||
page_models = [m for m in payload["data"] if isinstance(m, dict)]
|
||||
has_more = bool(payload.get("has_more"))
|
||||
elif isinstance(payload, list):
|
||||
page_models = [m for m in payload if isinstance(m, dict)]
|
||||
|
||||
for model in page_models:
|
||||
model.setdefault("api_format", normalized_api_format)
|
||||
|
||||
next_cursor = None
|
||||
if normalized_api_format.startswith("claude:") and isinstance(payload, dict):
|
||||
raw_last_id = payload.get("last_id")
|
||||
if has_more and isinstance(raw_last_id, str) and raw_last_id.strip():
|
||||
next_cursor = raw_last_id.strip()
|
||||
|
||||
return page_models, has_more, next_cursor
|
||||
|
||||
|
||||
def _extract_rust_error_message(result: Any) -> str:
|
||||
payload = getattr(result, "response_json", None)
|
||||
if isinstance(payload, dict):
|
||||
err = payload.get("error")
|
||||
if isinstance(err, dict):
|
||||
message = str(err.get("message") or "").strip()
|
||||
if message:
|
||||
return message[:500]
|
||||
message = str(payload.get("message") or "").strip()
|
||||
if message:
|
||||
return message[:500]
|
||||
try:
|
||||
return json.dumps(payload, ensure_ascii=False)[:500]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
body_bytes = getattr(result, "response_body_bytes", None)
|
||||
if body_bytes:
|
||||
try:
|
||||
return body_bytes.decode("utf-8", errors="replace")[:500]
|
||||
except Exception:
|
||||
return "(binary body)"
|
||||
return "(empty)"
|
||||
|
||||
|
||||
async def _try_rust_fetch_models_for_api_format(
|
||||
*,
|
||||
api_format: str,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
extra_headers: dict[str, str] | None,
|
||||
timeout: float,
|
||||
proxy_config: dict[str, Any] | None,
|
||||
) -> tuple[list[dict[str, Any]], str | None, bool] | None:
|
||||
from src.services.request.execution_runtime_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanBody,
|
||||
ExecutionPlanTimeouts,
|
||||
)
|
||||
from src.services.request.execution_runtime_client import (
|
||||
ExecutionRuntimeClient,
|
||||
ExecutionRuntimeClientError,
|
||||
)
|
||||
|
||||
if config.execution_runtime_backend != "rust":
|
||||
return None
|
||||
|
||||
try:
|
||||
proxy_snapshot = await _build_models_proxy_snapshot(proxy_config)
|
||||
headers = _build_model_fetch_headers(api_format, api_key, extra_headers)
|
||||
models: list[dict[str, Any]] = []
|
||||
next_cursor: str | None = None
|
||||
seen_ids: set[str] = set()
|
||||
|
||||
for _ in range(20):
|
||||
url = _build_model_fetch_url(api_format, base_url, api_key, next_cursor)
|
||||
result = await ExecutionRuntimeClient().execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id=f"model-fetch:{api_format}:{next_cursor or 'root'}",
|
||||
candidate_id=None,
|
||||
provider_name=str(api_format.split(":", 1)[0] or "unknown"),
|
||||
provider_id="",
|
||||
endpoint_id="",
|
||||
key_id="",
|
||||
method="GET",
|
||||
url=url,
|
||||
headers=dict(headers),
|
||||
body=ExecutionPlanBody(),
|
||||
stream=False,
|
||||
provider_api_format=api_format,
|
||||
client_api_format=api_format,
|
||||
model_name="models",
|
||||
proxy=proxy_snapshot,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=min(int(timeout * 1000), 30_000),
|
||||
read_ms=int(timeout * 1000),
|
||||
write_ms=int(timeout * 1000),
|
||||
pool_ms=min(int(timeout * 1000), 30_000),
|
||||
total_ms=int(timeout * 1000),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
if result.status_code != 200:
|
||||
error_body = _extract_rust_error_message(result)
|
||||
return [], f"HTTP {result.status_code}: {error_body}", False
|
||||
|
||||
payload = result.response_json
|
||||
page_models, has_more, next_after_id = _parse_model_fetch_response(api_format, payload)
|
||||
for model in page_models:
|
||||
model_id = model.get("id")
|
||||
if isinstance(model_id, str) and model_id and model_id in seen_ids:
|
||||
continue
|
||||
if isinstance(model_id, str) and model_id:
|
||||
seen_ids.add(model_id)
|
||||
models.append(model)
|
||||
|
||||
if not has_more or not next_after_id or next_after_id == next_cursor:
|
||||
return models, None, True
|
||||
next_cursor = next_after_id
|
||||
|
||||
return models, None, True
|
||||
except (ExecutionRuntimeClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
|
||||
logger.warning("Rust model fetch fallback for {}: {}", api_format, exc)
|
||||
return None
|
||||
except Exception as exc:
|
||||
logger.warning("Rust model fetch unexpected fallback for {}: {}", api_format, exc)
|
||||
return None
|
||||
|
||||
|
||||
async def fetch_models_from_endpoints(
|
||||
endpoint_configs: list[dict],
|
||||
timeout: float = 30.0,
|
||||
proxy_config: dict[str, Any] | None = None,
|
||||
) -> tuple[list[dict], list[str], bool]:
|
||||
"""
|
||||
从多个端点并发获取模型
|
||||
|
||||
Args:
|
||||
endpoint_configs: 端点配置列表,每个配置包含 api_key, base_url, api_format, extra_headers
|
||||
timeout: 请求超时时间(秒)
|
||||
proxy_config: 代理配置(可选),支持系统默认回退
|
||||
|
||||
Returns:
|
||||
(模型列表, 错误列表, 是否有成功)
|
||||
"""
|
||||
all_models: list[dict] = []
|
||||
errors: list[str] = []
|
||||
has_success = False
|
||||
semaphore = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)
|
||||
|
||||
async def fetch_one(config: dict) -> tuple[list, str | None, bool]:
|
||||
base_url = config["base_url"]
|
||||
if not base_url:
|
||||
return [], None, False
|
||||
base_url = base_url.rstrip("/")
|
||||
api_format = config["api_format"]
|
||||
api_key_value = config["api_key"]
|
||||
extra_headers = config.get("extra_headers")
|
||||
|
||||
try:
|
||||
async with semaphore:
|
||||
rust_result = await _try_rust_fetch_models_for_api_format(
|
||||
api_format=api_format,
|
||||
base_url=base_url,
|
||||
api_key=api_key_value,
|
||||
extra_headers=extra_headers,
|
||||
timeout=timeout,
|
||||
proxy_config=proxy_config,
|
||||
)
|
||||
if rust_result is None:
|
||||
return [], f"{api_format}: rust executor unavailable", False
|
||||
models, error, success = rust_result
|
||||
|
||||
for m in models:
|
||||
if "api_format" not in m:
|
||||
m["api_format"] = api_format
|
||||
|
||||
# 即使返回空列表,只要没有错误也算成功
|
||||
return models, error, success
|
||||
except httpx.TimeoutException:
|
||||
logger.warning("获取 {} 模型超时", api_format)
|
||||
return [], f"{api_format}: timeout", False
|
||||
except Exception:
|
||||
logger.exception("获取 {} 模型出错", api_format)
|
||||
return [], f"{api_format}: error", False
|
||||
|
||||
results = await asyncio.gather(*[fetch_one(c) for c in endpoint_configs])
|
||||
for models, error, success in results:
|
||||
all_models.extend(models)
|
||||
if error:
|
||||
errors.append(error)
|
||||
if success:
|
||||
has_success = True
|
||||
|
||||
return all_models, errors, has_success
|
||||
Reference in New Issue
Block a user