45 lines
1.8 KiB
Python
45 lines
1.8 KiB
Python
"""多 Provider 自动 failover（Hermes 融合：源自 error_classifier / credential_pool 思路）。
|
||
|
||
当主 Provider 的错误类型属于「换一家就能好」的类别（认证失败/限流/连接/超时/
|
||
模型不存在/余额不足）时，自动切到下一个启用的 Provider 重试一次；
|
||
模型不存在的用目标 Provider 的默认模型。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from sqlmodel import Session, select
|
||
|
||
from .models import ModelProvider
|
||
|
||
# 这些错误换 Provider 有意义；auth/balance 属于账号问题，换一家确实能好
|
||
FAILOVER_KINDS = frozenset({
|
||
"auth", "rate_limit", "connection", "timeout", "model_not_found", "balance",
|
||
})
|
||
|
||
|
||
def failover_candidates(session: Session, current_provider_id: int) -> list[ModelProvider]:
|
||
"""按 id 升序返回其他启用的 Provider（排除当前）。"""
|
||
rows = session.exec(select(ModelProvider).order_by(ModelProvider.id)).all()
|
||
return [p for p in rows
|
||
if p.id != current_provider_id and getattr(p, "enabled", False)]
|
||
|
||
|
||
def should_failover(kind: str) -> bool:
|
||
return kind in FAILOVER_KINDS
|
||
|
||
|
||
def failover_target(session: Session, current_provider_id: int,
|
||
current_model: str) -> tuple[ModelProvider, str] | None:
|
||
"""选一个 failover 目标：(Provider, 模型)。优先沿用同名模型，否则用其默认模型。"""
|
||
for provider in failover_candidates(session, current_provider_id):
|
||
models = list(provider.models or [])
|
||
if current_model in models:
|
||
return provider, current_model
|
||
if provider.default_model and provider.default_model in models:
|
||
return provider, provider.default_model
|
||
if models:
|
||
return provider, models[0]
|
||
return None
|
||
|
||
|
||
__all__ = ["FAILOVER_KINDS", "failover_candidates", "failover_target", "should_failover"]
|