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:
fawney19
2026-04-03 16:26:16 +08:00
parent 8f26e1a31f
commit 1d9c77522a
868 changed files with 1735 additions and 2433 deletions

View 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}")

View 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_availableis_available=NULL 视为可用,兼容历史数据)
"""
@staticmethod
def base_active_models(db: Session, eager_load: bool = False) -> Query:
"""
返回基础的可用模型查询(实体为 Model
已包含条件:
- Model.is_active = True
- Model.is_available = True 或 NULLNULL 视为可用)
- 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

File diff suppressed because it is too large Load Diff

View 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,
)
# ========== 阶段 2HTTP 请求(不持有数据库连接)==========
# 使用较短的超时时间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_keyOAuth: 加密 access_tokenAPI 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

View 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 KeyOAuth 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

View 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

View 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 # 命中的阶梯索引

View 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(),
)

View 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