mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +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:
20
_deprecated_py_src/services/provider/pool/__init__.py
Normal file
20
_deprecated_py_src/services/provider/pool/__init__.py
Normal file
@@ -0,0 +1,20 @@
|
||||
"""Generic Account Pool management for any Provider type.
|
||||
|
||||
Re-exports the main public API for convenience.
|
||||
"""
|
||||
|
||||
from src.services.provider.pool.config import (
|
||||
PoolConfig,
|
||||
ScoringWeights,
|
||||
UnschedulableRule,
|
||||
parse_pool_config,
|
||||
)
|
||||
from src.services.provider.pool.manager import PoolManager
|
||||
|
||||
__all__ = [
|
||||
"PoolConfig",
|
||||
"PoolManager",
|
||||
"ScoringWeights",
|
||||
"UnschedulableRule",
|
||||
"parse_pool_config",
|
||||
]
|
||||
594
_deprecated_py_src/services/provider/pool/account_state.py
Normal file
594
_deprecated_py_src/services/provider/pool/account_state.py
Normal file
@@ -0,0 +1,594 @@
|
||||
"""Pool account state helpers.
|
||||
|
||||
Provides a shared way to classify account-level hard-block states
|
||||
from upstream metadata and OAuth invalid reasons.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from src.services.provider_keys.quota_reader import get_quota_reader
|
||||
|
||||
OAUTH_ACCOUNT_BLOCK_PREFIX = "[ACCOUNT_BLOCK] "
|
||||
OAUTH_REFRESH_FAILED_PREFIX = "[REFRESH_FAILED] "
|
||||
OAUTH_EXPIRED_PREFIX = "[OAUTH_EXPIRED] "
|
||||
OAUTH_REQUEST_FAILED_PREFIX = "[REQUEST_FAILED] "
|
||||
|
||||
# -- 按原因细分的关键词组 --
|
||||
# 封禁类 (suspended / banned)
|
||||
_KEYWORDS_SUSPENDED: tuple[str, ...] = (
|
||||
"suspended",
|
||||
"account_block",
|
||||
"account blocked",
|
||||
"封禁",
|
||||
"封号",
|
||||
"被封",
|
||||
"账户已封禁",
|
||||
"账号异常",
|
||||
)
|
||||
|
||||
# 停用类 (disabled / deactivated)
|
||||
_KEYWORDS_DISABLED: tuple[str, ...] = (
|
||||
"account has been disabled",
|
||||
"account disabled",
|
||||
"account has been deactivated",
|
||||
"account_deactivated",
|
||||
"account deactivated",
|
||||
"organization has been disabled",
|
||||
"organization_disabled",
|
||||
"deactivated_workspace",
|
||||
"deactivated",
|
||||
"访问被禁止",
|
||||
"账户访问被禁止",
|
||||
)
|
||||
|
||||
_TOKEN_INVALID_KEYWORDS: tuple[str, ...] = (
|
||||
"authentication token has been invalidated",
|
||||
"token has been invalidated",
|
||||
"codex token 无效或已过期",
|
||||
)
|
||||
|
||||
# 需要验证类
|
||||
_KEYWORDS_VERIFICATION: tuple[str, ...] = (
|
||||
"validation_required",
|
||||
"verify your account",
|
||||
"需要验证",
|
||||
"验证账号",
|
||||
"验证身份",
|
||||
)
|
||||
|
||||
# 合并的完整列表(用于 is_account_level_block_reason 快速判断)
|
||||
ACCOUNT_BLOCK_REASON_KEYWORDS: tuple[str, ...] = (
|
||||
*_KEYWORDS_SUSPENDED,
|
||||
*_KEYWORDS_DISABLED,
|
||||
*_TOKEN_INVALID_KEYWORDS,
|
||||
*_KEYWORDS_VERIFICATION,
|
||||
)
|
||||
|
||||
AUTO_REMOVABLE_ACCOUNT_STATE_CODES: frozenset[str] = frozenset(
|
||||
{
|
||||
"account_banned",
|
||||
"account_suspended",
|
||||
"account_disabled",
|
||||
"workspace_deactivated",
|
||||
"account_forbidden",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _classify_block_reason(text: str) -> tuple[str, str]:
|
||||
"""Return (code, label) based on the oauth_invalid_reason text."""
|
||||
lowered = text.lower()
|
||||
if any(kw in lowered for kw in _TOKEN_INVALID_KEYWORDS):
|
||||
return "oauth_expired", "Token 失效"
|
||||
if any(kw in lowered for kw in _KEYWORDS_VERIFICATION):
|
||||
return "account_verification", "需要验证"
|
||||
if "deactivated_workspace" in lowered:
|
||||
return "workspace_deactivated", "工作区停用"
|
||||
if any(kw in lowered for kw in _KEYWORDS_DISABLED):
|
||||
return "account_disabled", "账号停用"
|
||||
if any(kw in lowered for kw in _KEYWORDS_SUSPENDED):
|
||||
return "account_suspended", "账号封禁"
|
||||
return "account_blocked", "账号异常"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PoolAccountState:
|
||||
"""Resolved account-level state for one key."""
|
||||
|
||||
blocked: bool
|
||||
code: str | None = None # account_banned / account_forbidden / account_blocked
|
||||
label: str | None = None
|
||||
reason: str | None = None
|
||||
source: str | None = None # metadata / oauth_invalid / oauth_refresh / oauth_request
|
||||
recoverable: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class OAuthStatusSnapshot:
|
||||
code: str = "none" # none / valid / expiring / expired / invalid / check_failed
|
||||
label: str | None = None
|
||||
reason: str | None = None
|
||||
expires_at: int | None = None
|
||||
invalid_at: int | None = None
|
||||
source: str | None = None
|
||||
requires_reauth: bool = False
|
||||
expiring_soon: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AccountStatusSnapshot:
|
||||
code: str = "ok"
|
||||
label: str | None = None
|
||||
reason: str | None = None
|
||||
blocked: bool = False
|
||||
source: str | None = None
|
||||
recoverable: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class QuotaStatusSnapshot:
|
||||
code: str = "unknown" # unknown / ok / exhausted
|
||||
label: str | None = None
|
||||
reason: str | None = None
|
||||
exhausted: bool = False
|
||||
usage_ratio: float | None = None
|
||||
updated_at: int | None = None
|
||||
reset_seconds: float | None = None
|
||||
plan_type: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProviderKeyStatusSnapshot:
|
||||
oauth: OAuthStatusSnapshot
|
||||
account: AccountStatusSnapshot
|
||||
quota: QuotaStatusSnapshot
|
||||
|
||||
|
||||
def _is_truthy_flag(value: Any) -> bool:
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, (int, float)):
|
||||
return value != 0
|
||||
if isinstance(value, str):
|
||||
normalized = value.strip().lower()
|
||||
return normalized in {"1", "true", "yes", "y"}
|
||||
return False
|
||||
|
||||
|
||||
def _clean_text(value: Any) -> str | None:
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
text = value.strip()
|
||||
return text or None
|
||||
|
||||
|
||||
def _extract_reason(source: dict[str, Any] | None, *fields: str) -> str | None:
|
||||
if not isinstance(source, dict):
|
||||
return None
|
||||
for field in fields:
|
||||
text = _clean_text(source.get(field))
|
||||
if text:
|
||||
return text
|
||||
return None
|
||||
|
||||
|
||||
def _is_workspace_deactivated_reason(reason: str | None) -> bool:
|
||||
text = _clean_text(reason)
|
||||
return bool(text and "deactivated_workspace" in text.lower())
|
||||
|
||||
|
||||
_TAGGED_REASON_PATTERN = re.compile(
|
||||
r"(?:^|\n)\[(?P<tag>[A-Z_]+)\]\s*(?P<detail>.*?)(?=\n\[[A-Z_]+\]|\Z)",
|
||||
re.S,
|
||||
)
|
||||
|
||||
|
||||
def _extract_tagged_reason_sections(reason: str | None) -> dict[str, str]:
|
||||
text = _clean_text(reason)
|
||||
if not text:
|
||||
return {}
|
||||
sections: dict[str, str] = {}
|
||||
for match in _TAGGED_REASON_PATTERN.finditer(text):
|
||||
tag = str(match.group("tag") or "").strip().upper()
|
||||
if not tag or tag in sections:
|
||||
continue
|
||||
detail = str(match.group("detail") or "").strip()
|
||||
sections[tag] = detail
|
||||
return sections
|
||||
|
||||
|
||||
def _resolve_from_metadata(
|
||||
provider_type: str | None,
|
||||
upstream_metadata: Any,
|
||||
) -> PoolAccountState | None:
|
||||
if not isinstance(upstream_metadata, dict):
|
||||
return None
|
||||
|
||||
normalized_provider = str(provider_type or "").strip().lower()
|
||||
provider_bucket: dict[str, Any] | None = None
|
||||
if normalized_provider:
|
||||
maybe_bucket = upstream_metadata.get(normalized_provider)
|
||||
if isinstance(maybe_bucket, dict):
|
||||
provider_bucket = maybe_bucket
|
||||
|
||||
quota_block = get_quota_reader(normalized_provider, upstream_metadata).account_block()
|
||||
if quota_block.blocked:
|
||||
return PoolAccountState(
|
||||
blocked=True,
|
||||
code=quota_block.code,
|
||||
label=quota_block.label,
|
||||
reason=quota_block.reason,
|
||||
source="metadata",
|
||||
)
|
||||
|
||||
for source in (provider_bucket, upstream_metadata):
|
||||
if not isinstance(source, dict):
|
||||
continue
|
||||
if _is_truthy_flag(source.get("is_banned")):
|
||||
reason = _extract_reason(source, "ban_reason", "forbidden_reason", "reason", "message")
|
||||
return PoolAccountState(
|
||||
blocked=True,
|
||||
code="account_banned",
|
||||
label="账号封禁",
|
||||
reason=reason or "账号已封禁",
|
||||
source="metadata",
|
||||
)
|
||||
if _is_truthy_flag(source.get("is_forbidden")) or _is_truthy_flag(
|
||||
source.get("account_disabled")
|
||||
):
|
||||
reason = _extract_reason(source, "forbidden_reason", "ban_reason", "reason", "message")
|
||||
if _is_workspace_deactivated_reason(reason):
|
||||
return PoolAccountState(
|
||||
blocked=True,
|
||||
code="workspace_deactivated",
|
||||
label="工作区停用",
|
||||
reason=reason or "工作区已停用",
|
||||
source="metadata",
|
||||
)
|
||||
return PoolAccountState(
|
||||
blocked=True,
|
||||
code="account_forbidden",
|
||||
label="访问受限",
|
||||
reason=reason or "账号访问受限",
|
||||
source="metadata",
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_from_oauth_invalid_reason(reason: str | None) -> PoolAccountState | None:
|
||||
text = _clean_text(reason)
|
||||
if not text:
|
||||
return None
|
||||
|
||||
if text.startswith(OAUTH_ACCOUNT_BLOCK_PREFIX):
|
||||
cleaned = text[len(OAUTH_ACCOUNT_BLOCK_PREFIX) :].strip()
|
||||
code, label = (
|
||||
_classify_block_reason(cleaned) if cleaned else ("account_blocked", "账号异常")
|
||||
)
|
||||
return PoolAccountState(
|
||||
blocked=True,
|
||||
code=code,
|
||||
label=label,
|
||||
reason=cleaned or "账号异常",
|
||||
source="oauth_invalid",
|
||||
)
|
||||
|
||||
if text.startswith(OAUTH_EXPIRED_PREFIX):
|
||||
cleaned = text[len(OAUTH_EXPIRED_PREFIX) :].strip()
|
||||
return PoolAccountState(
|
||||
blocked=True,
|
||||
code="oauth_expired",
|
||||
label="Token 失效",
|
||||
reason=cleaned or "OAuth Token 已过期且无法续期",
|
||||
source="oauth_invalid",
|
||||
recoverable=True,
|
||||
)
|
||||
|
||||
if text.startswith(OAUTH_REFRESH_FAILED_PREFIX):
|
||||
cleaned = text[len(OAUTH_REFRESH_FAILED_PREFIX) :].strip()
|
||||
return PoolAccountState(
|
||||
blocked=False,
|
||||
code="oauth_refresh_failed",
|
||||
label="续期失败",
|
||||
reason=cleaned or "OAuth Token 续期失败",
|
||||
source="oauth_refresh",
|
||||
recoverable=True,
|
||||
)
|
||||
|
||||
if text.startswith(OAUTH_REQUEST_FAILED_PREFIX):
|
||||
cleaned = text[len(OAUTH_REQUEST_FAILED_PREFIX) :].strip()
|
||||
return PoolAccountState(
|
||||
blocked=False,
|
||||
code="oauth_request_failed",
|
||||
label="请求失败",
|
||||
reason=cleaned or "账号状态检查失败",
|
||||
source="oauth_request",
|
||||
recoverable=True,
|
||||
)
|
||||
|
||||
if text.startswith("["):
|
||||
return None
|
||||
|
||||
lowered = text.lower()
|
||||
if any(keyword in lowered for keyword in ACCOUNT_BLOCK_REASON_KEYWORDS):
|
||||
code, label = _classify_block_reason(text)
|
||||
return PoolAccountState(
|
||||
blocked=True,
|
||||
code=code,
|
||||
label=label,
|
||||
reason=text,
|
||||
source="oauth_invalid",
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def resolve_account_status_snapshot(
|
||||
*,
|
||||
provider_type: str | None,
|
||||
upstream_metadata: Any,
|
||||
oauth_invalid_reason: str | None,
|
||||
) -> AccountStatusSnapshot:
|
||||
from_metadata = _resolve_from_metadata(provider_type, upstream_metadata)
|
||||
if from_metadata is not None:
|
||||
return AccountStatusSnapshot(
|
||||
code=from_metadata.code or "ok",
|
||||
label=from_metadata.label,
|
||||
reason=from_metadata.reason,
|
||||
blocked=from_metadata.blocked,
|
||||
source=from_metadata.source,
|
||||
recoverable=from_metadata.recoverable,
|
||||
)
|
||||
|
||||
text = _clean_text(oauth_invalid_reason)
|
||||
if not text:
|
||||
return AccountStatusSnapshot()
|
||||
|
||||
tagged_sections = _extract_tagged_reason_sections(text)
|
||||
if "ACCOUNT_BLOCK" in tagged_sections:
|
||||
cleaned = tagged_sections["ACCOUNT_BLOCK"]
|
||||
code, label = (
|
||||
_classify_block_reason(cleaned) if cleaned else ("account_blocked", "账号异常")
|
||||
)
|
||||
return AccountStatusSnapshot(
|
||||
code=code,
|
||||
label=label,
|
||||
reason=cleaned or "账号异常",
|
||||
blocked=True,
|
||||
source="oauth_invalid",
|
||||
)
|
||||
|
||||
if text.startswith("["):
|
||||
return AccountStatusSnapshot()
|
||||
|
||||
lowered = text.lower()
|
||||
if any(keyword in lowered for keyword in ACCOUNT_BLOCK_REASON_KEYWORDS):
|
||||
code, label = _classify_block_reason(text)
|
||||
return AccountStatusSnapshot(
|
||||
code=code,
|
||||
label=label,
|
||||
reason=text,
|
||||
blocked=True,
|
||||
source="oauth_invalid",
|
||||
)
|
||||
|
||||
return AccountStatusSnapshot()
|
||||
|
||||
|
||||
def resolve_oauth_status_snapshot(
|
||||
*,
|
||||
auth_type: str | None,
|
||||
oauth_expires_at: int | None,
|
||||
oauth_invalid_at: int | None,
|
||||
oauth_invalid_reason: str | None,
|
||||
now_ts: int | None = None,
|
||||
) -> OAuthStatusSnapshot:
|
||||
if str(auth_type or "").strip().lower() != "oauth":
|
||||
return OAuthStatusSnapshot()
|
||||
|
||||
now = int(now_ts if now_ts is not None else time.time())
|
||||
invalid_at = int(oauth_invalid_at) if isinstance(oauth_invalid_at, int) else None
|
||||
tagged_sections = _extract_tagged_reason_sections(oauth_invalid_reason)
|
||||
raw_reason = _clean_text(oauth_invalid_reason)
|
||||
|
||||
expired_reason = tagged_sections.get("OAUTH_EXPIRED")
|
||||
if expired_reason:
|
||||
return OAuthStatusSnapshot(
|
||||
code="invalid",
|
||||
label="已失效",
|
||||
reason=expired_reason,
|
||||
invalid_at=invalid_at,
|
||||
expires_at=oauth_expires_at,
|
||||
source="oauth_invalid",
|
||||
requires_reauth=True,
|
||||
)
|
||||
|
||||
refresh_failed_reason = tagged_sections.get("REFRESH_FAILED")
|
||||
if refresh_failed_reason:
|
||||
return OAuthStatusSnapshot(
|
||||
code="invalid",
|
||||
label="已失效",
|
||||
reason=refresh_failed_reason,
|
||||
invalid_at=invalid_at,
|
||||
expires_at=oauth_expires_at,
|
||||
source="oauth_refresh",
|
||||
requires_reauth=True,
|
||||
)
|
||||
|
||||
request_failed_reason = tagged_sections.get("REQUEST_FAILED")
|
||||
if request_failed_reason:
|
||||
return OAuthStatusSnapshot(
|
||||
code="check_failed",
|
||||
label="检查失败",
|
||||
reason=request_failed_reason,
|
||||
expires_at=oauth_expires_at,
|
||||
source="oauth_request",
|
||||
)
|
||||
|
||||
account_snapshot = resolve_account_status_snapshot(
|
||||
provider_type=None,
|
||||
upstream_metadata=None,
|
||||
oauth_invalid_reason=raw_reason,
|
||||
)
|
||||
if account_snapshot.blocked:
|
||||
if oauth_expires_at is None:
|
||||
return OAuthStatusSnapshot()
|
||||
elif raw_reason or invalid_at is not None:
|
||||
return OAuthStatusSnapshot(
|
||||
code="invalid",
|
||||
label="已失效",
|
||||
reason=raw_reason,
|
||||
invalid_at=invalid_at,
|
||||
expires_at=oauth_expires_at,
|
||||
source="oauth_invalid",
|
||||
requires_reauth=True,
|
||||
)
|
||||
|
||||
expires_at = int(oauth_expires_at) if isinstance(oauth_expires_at, int) else None
|
||||
if expires_at is None:
|
||||
return OAuthStatusSnapshot()
|
||||
if expires_at <= now:
|
||||
return OAuthStatusSnapshot(
|
||||
code="expired",
|
||||
label="已过期",
|
||||
reason="Token 已过期,请重新授权",
|
||||
expires_at=expires_at,
|
||||
source="expires_at",
|
||||
requires_reauth=True,
|
||||
)
|
||||
expiring_soon = (expires_at - now) < 24 * 3600
|
||||
return OAuthStatusSnapshot(
|
||||
code="expiring" if expiring_soon else "valid",
|
||||
label="即将过期" if expiring_soon else "有效",
|
||||
expires_at=expires_at,
|
||||
source="expires_at",
|
||||
expiring_soon=expiring_soon,
|
||||
)
|
||||
|
||||
|
||||
def resolve_quota_status_snapshot(
|
||||
*,
|
||||
provider_type: str | None,
|
||||
upstream_metadata: Any,
|
||||
) -> QuotaStatusSnapshot:
|
||||
normalized_provider = str(provider_type or "").strip().lower()
|
||||
reader = get_quota_reader(normalized_provider, upstream_metadata)
|
||||
quota_state = reader.is_exhausted()
|
||||
usage_ratio = reader.usage_ratio()
|
||||
updated_at = reader.updated_at()
|
||||
reset_seconds = reader.reset_seconds()
|
||||
plan_type = reader.plan_type()
|
||||
|
||||
if quota_state.exhausted:
|
||||
return QuotaStatusSnapshot(
|
||||
code="exhausted",
|
||||
label="额度耗尽",
|
||||
reason=quota_state.reason,
|
||||
exhausted=True,
|
||||
usage_ratio=usage_ratio,
|
||||
updated_at=updated_at,
|
||||
reset_seconds=reset_seconds,
|
||||
plan_type=plan_type,
|
||||
)
|
||||
|
||||
if any(value is not None for value in (usage_ratio, updated_at, reset_seconds, plan_type)):
|
||||
return QuotaStatusSnapshot(
|
||||
code="ok",
|
||||
exhausted=False,
|
||||
usage_ratio=usage_ratio,
|
||||
updated_at=updated_at,
|
||||
reset_seconds=reset_seconds,
|
||||
plan_type=plan_type,
|
||||
)
|
||||
|
||||
return QuotaStatusSnapshot()
|
||||
|
||||
|
||||
def build_provider_key_status_snapshot(
|
||||
*,
|
||||
auth_type: str | None,
|
||||
oauth_expires_at: int | None,
|
||||
oauth_invalid_at: int | None,
|
||||
oauth_invalid_reason: str | None,
|
||||
provider_type: str | None,
|
||||
upstream_metadata: Any,
|
||||
now_ts: int | None = None,
|
||||
) -> ProviderKeyStatusSnapshot:
|
||||
account = resolve_account_status_snapshot(
|
||||
provider_type=provider_type,
|
||||
upstream_metadata=upstream_metadata,
|
||||
oauth_invalid_reason=oauth_invalid_reason,
|
||||
)
|
||||
oauth = resolve_oauth_status_snapshot(
|
||||
auth_type=auth_type,
|
||||
oauth_expires_at=oauth_expires_at,
|
||||
oauth_invalid_at=oauth_invalid_at,
|
||||
oauth_invalid_reason=oauth_invalid_reason,
|
||||
now_ts=now_ts,
|
||||
)
|
||||
quota = resolve_quota_status_snapshot(
|
||||
provider_type=provider_type,
|
||||
upstream_metadata=upstream_metadata,
|
||||
)
|
||||
return ProviderKeyStatusSnapshot(oauth=oauth, account=account, quota=quota)
|
||||
|
||||
|
||||
def resolve_pool_account_state(
|
||||
*,
|
||||
provider_type: str | None,
|
||||
upstream_metadata: Any,
|
||||
oauth_invalid_reason: str | None,
|
||||
) -> PoolAccountState:
|
||||
"""Resolve account-level hard-block state for pool scheduling."""
|
||||
|
||||
from_metadata = _resolve_from_metadata(provider_type, upstream_metadata)
|
||||
if from_metadata is not None:
|
||||
return from_metadata
|
||||
|
||||
from_oauth = _resolve_from_oauth_invalid_reason(oauth_invalid_reason)
|
||||
if from_oauth is not None:
|
||||
return from_oauth
|
||||
|
||||
return PoolAccountState(blocked=False)
|
||||
|
||||
|
||||
def should_auto_remove_account_state(state: PoolAccountState) -> bool:
|
||||
"""Whether a resolved account state is safe to auto-remove.
|
||||
|
||||
Auto-removal is limited to hard, non-recoverable account abnormalities.
|
||||
Pure token failures (`oauth_expired`, `oauth_refresh_failed`) and
|
||||
softer/manual-recoverable states like `account_verification` are excluded.
|
||||
"""
|
||||
|
||||
return bool(
|
||||
state.blocked
|
||||
and not state.recoverable
|
||||
and str(state.code or "").strip().lower() in AUTO_REMOVABLE_ACCOUNT_STATE_CODES
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ACCOUNT_BLOCK_REASON_KEYWORDS",
|
||||
"AUTO_REMOVABLE_ACCOUNT_STATE_CODES",
|
||||
"AccountStatusSnapshot",
|
||||
"OAUTH_ACCOUNT_BLOCK_PREFIX",
|
||||
"OAUTH_EXPIRED_PREFIX",
|
||||
"OAUTH_REFRESH_FAILED_PREFIX",
|
||||
"OAUTH_REQUEST_FAILED_PREFIX",
|
||||
"OAuthStatusSnapshot",
|
||||
"PoolAccountState",
|
||||
"ProviderKeyStatusSnapshot",
|
||||
"QuotaStatusSnapshot",
|
||||
"build_provider_key_status_snapshot",
|
||||
"resolve_account_status_snapshot",
|
||||
"resolve_oauth_status_snapshot",
|
||||
"resolve_pool_account_state",
|
||||
"resolve_quota_status_snapshot",
|
||||
"should_auto_remove_account_state",
|
||||
]
|
||||
341
_deprecated_py_src/services/provider/pool/config.py
Normal file
341
_deprecated_py_src/services/provider/pool/config.py
Normal file
@@ -0,0 +1,341 @@
|
||||
"""Account Pool configuration (provider-agnostic)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.pool.dimensions import get_preset_dimension, get_preset_names
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ScoringWeights:
|
||||
"""Weights used by multi-score scheduling."""
|
||||
|
||||
lru: float = 0.3
|
||||
latency: float = 0.25
|
||||
health: float = 0.2
|
||||
cost_remaining: float = 0.25
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SchedulingPreset:
|
||||
"""Single scheduling preset item with enable/disable and optional sub-config."""
|
||||
|
||||
preset: str
|
||||
enabled: bool = True
|
||||
mode: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UnschedulableRule:
|
||||
"""Keyword-based temporary unschedule rule."""
|
||||
|
||||
keyword: str
|
||||
duration_minutes: int = 5
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PoolConfig:
|
||||
"""Parsed pool configuration for any Provider.
|
||||
|
||||
All transient state lives in Redis; this dataclass only holds
|
||||
the *configuration* that controls pool behaviour.
|
||||
"""
|
||||
|
||||
# -- Sticky Session -------------------------------------------------------
|
||||
sticky_session_ttl_seconds: int = 3600 # 1 hour
|
||||
# Key 优先模式下号池整体优先级(None 时回退 provider_priority)
|
||||
global_priority: int | None = None
|
||||
|
||||
# -- Load-Aware Selection -------------------------------------------------
|
||||
load_threshold_percent: int = 80
|
||||
|
||||
# -- Scheduling (unified preset list) -------------------------------------
|
||||
scheduling_presets: tuple[SchedulingPreset, ...] = (
|
||||
SchedulingPreset(preset="cache_affinity", enabled=True),
|
||||
)
|
||||
# Derived from scheduling_presets at parse time (backward compat for consumers)
|
||||
lru_enabled: bool = True
|
||||
scheduling_mode: str = "lru" # lru | multi_score
|
||||
|
||||
scoring_weights: ScoringWeights = field(default_factory=ScoringWeights)
|
||||
latency_window_seconds: int = 3600
|
||||
latency_sample_limit: int = 50
|
||||
|
||||
# -- Rolling-Window Cost Tracking -----------------------------------------
|
||||
cost_window_seconds: int = 18000 # 5 hours
|
||||
cost_limit_per_key_tokens: int | None = None # None = unlimited
|
||||
cost_soft_threshold_percent: int = 80
|
||||
|
||||
# -- Cooldown Defaults ----------------------------------------------------
|
||||
rate_limit_cooldown_seconds: int = 300 # 429
|
||||
overload_cooldown_seconds: int = 30 # 529
|
||||
|
||||
# -- OAuth Proactive Refresh ----------------------------------------------
|
||||
proactive_refresh_seconds: int = 180 # 3 minutes before expiry
|
||||
|
||||
# -- Health Policy --------------------------------------------------------
|
||||
health_policy_enabled: bool = True
|
||||
|
||||
# -- Temporary Unschedulable Rules ----------------------------------------
|
||||
unschedulable_rules: list[UnschedulableRule] = field(default_factory=list)
|
||||
|
||||
# -- Batch Operations -----------------------------------------------------
|
||||
batch_concurrency: int = 8
|
||||
|
||||
# -- Quota Probing --------------------------------------------------------
|
||||
probing_enabled: bool = False
|
||||
probing_interval_minutes: int = 10
|
||||
auto_remove_banned_keys: bool = False
|
||||
|
||||
# -- Stream Timeout Auto-Pause --------------------------------------------
|
||||
stream_timeout_threshold: int = 3 # N timeouts within window trigger cooldown
|
||||
stream_timeout_window_seconds: int = 1800 # 30 min counting window
|
||||
stream_timeout_cooldown_seconds: int = 300 # 5 min cooldown
|
||||
|
||||
# -- Pluggable Strategies -------------------------------------------------
|
||||
strategies: tuple[str, ...] = ()
|
||||
|
||||
|
||||
def parse_pool_config(provider_config: Any) -> PoolConfig | None:
|
||||
"""Parse PoolConfig from ``Provider.config``.
|
||||
|
||||
Only looks for the explicit ``pool_advanced`` key. Returns ``None``
|
||||
when the provider has no pool section configured, meaning the caller
|
||||
should use the normal (non-pool) scheduling path.
|
||||
"""
|
||||
config_dict = provider_config if isinstance(provider_config, dict) else {}
|
||||
|
||||
raw_advanced = config_dict.get("pool_advanced")
|
||||
if raw_advanced is None:
|
||||
return None
|
||||
|
||||
if not isinstance(raw_advanced, dict):
|
||||
# Could be a pre-validated Pydantic model; grab its dict.
|
||||
try:
|
||||
raw_advanced = raw_advanced.model_dump() # type: ignore[union-attr]
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"PoolConfig: advanced config type invalid ({}), falling back to defaults",
|
||||
type(raw_advanced).__name__,
|
||||
)
|
||||
return PoolConfig()
|
||||
|
||||
rules: list[UnschedulableRule] = []
|
||||
raw_rules = raw_advanced.get("unschedulable_rules")
|
||||
if isinstance(raw_rules, list):
|
||||
for r in raw_rules:
|
||||
if isinstance(r, dict) and isinstance(r.get("keyword"), str):
|
||||
rules.append(
|
||||
UnschedulableRule(
|
||||
keyword=r["keyword"],
|
||||
duration_minutes=int(r.get("duration_minutes", 5)),
|
||||
)
|
||||
)
|
||||
|
||||
def _int_or(key: str, default: int) -> int:
|
||||
v = raw_advanced.get(key)
|
||||
if v is None:
|
||||
return default
|
||||
try:
|
||||
return int(v)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
def _bool_or(key: str, default: bool) -> bool:
|
||||
v = raw_advanced.get(key)
|
||||
if v is None:
|
||||
return default
|
||||
return bool(v)
|
||||
|
||||
def _opt_int(key: str) -> int | None:
|
||||
v = raw_advanced.get(key)
|
||||
if v is None:
|
||||
return None
|
||||
try:
|
||||
return int(v)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
scoring_weights = _parse_scoring_weights(raw_advanced.get("scoring_weights"))
|
||||
|
||||
# Parse scheduling presets (new object-list format or legacy string-list)
|
||||
presets = _parse_scheduling_presets_v2(
|
||||
raw_advanced.get("scheduling_presets"),
|
||||
legacy_mode=raw_advanced.get("scheduling_mode"),
|
||||
legacy_lru=raw_advanced.get("lru_enabled"),
|
||||
)
|
||||
|
||||
# Derive scheduling_mode and lru_enabled from the presets list
|
||||
enabled = [p for p in presets if p.enabled]
|
||||
lru_enabled = any(p.preset == "lru" for p in enabled)
|
||||
non_lru_enabled = [p for p in enabled if p.preset != "lru"]
|
||||
scheduling_mode = "multi_score" if non_lru_enabled else "lru"
|
||||
|
||||
strategies = list(_parse_strategies(raw_advanced.get("strategies")))
|
||||
if scheduling_mode == "multi_score" and "multi_score" not in strategies:
|
||||
strategies.append("multi_score")
|
||||
|
||||
return PoolConfig(
|
||||
sticky_session_ttl_seconds=_int_or("sticky_session_ttl_seconds", 3600),
|
||||
global_priority=_opt_int("global_priority"),
|
||||
load_threshold_percent=_int_or("load_threshold_percent", 80),
|
||||
scheduling_presets=presets,
|
||||
lru_enabled=lru_enabled,
|
||||
scheduling_mode=scheduling_mode,
|
||||
scoring_weights=scoring_weights,
|
||||
latency_window_seconds=_int_or("latency_window_seconds", 3600),
|
||||
latency_sample_limit=_int_or("latency_sample_limit", 50),
|
||||
cost_window_seconds=_int_or("cost_window_seconds", 18000),
|
||||
cost_limit_per_key_tokens=_opt_int("cost_limit_per_key_tokens"),
|
||||
cost_soft_threshold_percent=_int_or("cost_soft_threshold_percent", 80),
|
||||
rate_limit_cooldown_seconds=_int_or("rate_limit_cooldown_seconds", 300),
|
||||
overload_cooldown_seconds=_int_or("overload_cooldown_seconds", 30),
|
||||
proactive_refresh_seconds=_int_or("proactive_refresh_seconds", 180),
|
||||
health_policy_enabled=_bool_or("health_policy_enabled", True),
|
||||
unschedulable_rules=rules,
|
||||
batch_concurrency=max(1, min(_int_or("batch_concurrency", 8), 32)),
|
||||
probing_enabled=_bool_or("probing_enabled", False),
|
||||
probing_interval_minutes=max(1, min(_int_or("probing_interval_minutes", 10), 1440)),
|
||||
auto_remove_banned_keys=_bool_or("auto_remove_banned_keys", False),
|
||||
stream_timeout_threshold=_int_or("stream_timeout_threshold", 3),
|
||||
stream_timeout_window_seconds=_int_or("stream_timeout_window_seconds", 1800),
|
||||
stream_timeout_cooldown_seconds=_int_or("stream_timeout_cooldown_seconds", 300),
|
||||
strategies=tuple(strategies),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Internal parsers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _allowed_preset_names() -> set[str]:
|
||||
return get_preset_names() | {"lru"}
|
||||
|
||||
|
||||
def _get_preset_mode_meta(preset_name: str) -> tuple[tuple[str, ...], str | None]:
|
||||
dim = get_preset_dimension(preset_name)
|
||||
if dim is None or not dim.modes:
|
||||
return (), None
|
||||
|
||||
modes = tuple(str(mode).strip().lower() for mode in dim.modes if str(mode).strip())
|
||||
if not modes:
|
||||
return (), None
|
||||
|
||||
raw_default = str(dim.default_mode or "").strip().lower()
|
||||
default_mode = raw_default if raw_default in modes else modes[0]
|
||||
return modes, default_mode
|
||||
|
||||
|
||||
def _parse_strategies(raw: Any) -> tuple[str, ...]:
|
||||
"""Parse strategy names from config (list[str] -> tuple[str, ...])."""
|
||||
if not isinstance(raw, list):
|
||||
return ()
|
||||
return tuple(str(s) for s in raw if isinstance(s, str) and s)
|
||||
|
||||
|
||||
def _parse_scoring_weights(raw: Any) -> ScoringWeights:
|
||||
"""Parse scoring weights with graceful fallback."""
|
||||
if not isinstance(raw, dict):
|
||||
return ScoringWeights()
|
||||
|
||||
def _float_or(value: Any, default: float) -> float:
|
||||
try:
|
||||
parsed = float(value)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
return max(0.0, min(parsed, 1.0))
|
||||
|
||||
return ScoringWeights(
|
||||
lru=_float_or(raw.get("lru"), 0.3),
|
||||
latency=_float_or(raw.get("latency"), 0.25),
|
||||
health=_float_or(raw.get("health"), 0.2),
|
||||
cost_remaining=_float_or(raw.get("cost_remaining"), 0.25),
|
||||
)
|
||||
|
||||
|
||||
def _parse_scheduling_presets_v2(
|
||||
raw: Any,
|
||||
*,
|
||||
legacy_mode: Any = None,
|
||||
legacy_lru: Any = None,
|
||||
) -> tuple[SchedulingPreset, ...]:
|
||||
"""Parse scheduling presets, supporting both new and legacy formats.
|
||||
|
||||
New format::
|
||||
|
||||
[{"preset": "lru", "enabled": true},
|
||||
{"preset": "free_team_first", "enabled": true, "mode": "free_only"},
|
||||
...]
|
||||
|
||||
Legacy format::
|
||||
|
||||
["free_team_first", "recent_refresh"] (with separate scheduling_mode / lru_enabled)
|
||||
"""
|
||||
if isinstance(raw, list) and raw:
|
||||
first = raw[0]
|
||||
if isinstance(first, dict):
|
||||
return _parse_preset_object_list(raw)
|
||||
if isinstance(first, str):
|
||||
return _convert_legacy_string_list(raw, legacy_mode, legacy_lru)
|
||||
|
||||
# No presets at all: derive from legacy fields
|
||||
return _build_from_legacy_fields(legacy_mode, legacy_lru)
|
||||
|
||||
|
||||
def _parse_preset_object_list(raw: list[Any]) -> tuple[SchedulingPreset, ...]:
|
||||
"""Parse new-format object list into SchedulingPreset tuple."""
|
||||
allowed = _allowed_preset_names()
|
||||
ordered: list[SchedulingPreset] = []
|
||||
seen: set[str] = set()
|
||||
for item in raw:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
name = str(item.get("preset", "")).strip().lower()
|
||||
if name not in allowed or name in seen:
|
||||
continue
|
||||
seen.add(name)
|
||||
enabled = bool(item.get("enabled", True))
|
||||
mode: str | None = None
|
||||
modes, default_mode = _get_preset_mode_meta(name)
|
||||
if modes:
|
||||
raw_mode = str(item.get("mode", default_mode) or "").strip().lower()
|
||||
mode = raw_mode if raw_mode in modes else default_mode
|
||||
ordered.append(SchedulingPreset(preset=name, enabled=enabled, mode=mode))
|
||||
return tuple(ordered) if ordered else (SchedulingPreset(preset="lru", enabled=True),)
|
||||
|
||||
|
||||
def _convert_legacy_string_list(
|
||||
raw: list[Any],
|
||||
legacy_mode: Any,
|
||||
legacy_lru: Any,
|
||||
) -> tuple[SchedulingPreset, ...]:
|
||||
"""Convert legacy string list + mode/lru fields to new format."""
|
||||
lru_enabled = legacy_lru if isinstance(legacy_lru, bool) else True
|
||||
|
||||
allowed_non_lru = _allowed_preset_names() - {"lru"}
|
||||
items: list[SchedulingPreset] = [SchedulingPreset(preset="lru", enabled=lru_enabled)]
|
||||
seen: set[str] = {"lru"}
|
||||
for p in raw:
|
||||
if not isinstance(p, str):
|
||||
continue
|
||||
name = p.strip().lower()
|
||||
if name not in allowed_non_lru or name in seen:
|
||||
continue
|
||||
seen.add(name)
|
||||
items.append(SchedulingPreset(preset=name, enabled=True))
|
||||
return tuple(items)
|
||||
|
||||
|
||||
def _build_from_legacy_fields(legacy_mode: Any, legacy_lru: Any) -> tuple[SchedulingPreset, ...]:
|
||||
"""Build presets from legacy scheduling_mode / lru_enabled only."""
|
||||
# Explicit lru_enabled=True -> LRU; explicit lru_enabled=False -> cache_affinity.
|
||||
# No legacy fields at all -> default to cache_affinity.
|
||||
if isinstance(legacy_lru, bool):
|
||||
if legacy_lru:
|
||||
return (SchedulingPreset(preset="lru", enabled=True),)
|
||||
return (SchedulingPreset(preset="cache_affinity", enabled=True),)
|
||||
return (SchedulingPreset(preset="cache_affinity", enabled=True),)
|
||||
66
_deprecated_py_src/services/provider/pool/cost_tracker.py
Normal file
66
_deprecated_py_src/services/provider/pool/cost_tracker.py
Normal file
@@ -0,0 +1,66 @@
|
||||
"""Rolling-window cost tracking for the Account Pool.
|
||||
|
||||
Each key has a configurable token budget per rolling window (e.g. 5 hours).
|
||||
When the budget is exhausted the key is marked as unschedulable by the pool
|
||||
manager. A "soft threshold" (default 80 %) causes the pool to *prefer*
|
||||
other keys but still allows traffic if no alternatives exist.
|
||||
|
||||
All state is stored in Redis sorted sets via :mod:`redis_ops`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from src.services.provider.pool import redis_ops
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
|
||||
|
||||
async def record_usage(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
tokens: int,
|
||||
config: PoolConfig,
|
||||
) -> None:
|
||||
"""Record *tokens* used by *key_id* in the rolling cost window."""
|
||||
if tokens <= 0:
|
||||
return
|
||||
if config.cost_limit_per_key_tokens is None:
|
||||
return # cost tracking disabled
|
||||
await redis_ops.add_cost_entry(provider_id, key_id, tokens, config.cost_window_seconds)
|
||||
|
||||
|
||||
async def get_window_usage(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
config: PoolConfig,
|
||||
) -> int:
|
||||
"""Return total tokens used by *key_id* within the current window."""
|
||||
return await redis_ops.get_cost_window_total(provider_id, key_id, config.cost_window_seconds)
|
||||
|
||||
|
||||
async def is_at_limit(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
config: PoolConfig,
|
||||
) -> bool:
|
||||
"""Return ``True`` if the key has exhausted its budget."""
|
||||
if config.cost_limit_per_key_tokens is None:
|
||||
return False
|
||||
total = await get_window_usage(provider_id, key_id, config)
|
||||
return total >= config.cost_limit_per_key_tokens
|
||||
|
||||
|
||||
async def is_approaching_limit(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
config: PoolConfig,
|
||||
) -> bool:
|
||||
"""Return ``True`` if the key is above the soft threshold."""
|
||||
if config.cost_limit_per_key_tokens is None:
|
||||
return False
|
||||
total = await get_window_usage(provider_id, key_id, config)
|
||||
threshold = config.cost_limit_per_key_tokens * config.cost_soft_threshold_percent / 100
|
||||
return total >= threshold
|
||||
@@ -0,0 +1,39 @@
|
||||
"""Pool scheduling preset dimensions.
|
||||
|
||||
Importing this package registers all built-in preset dimensions.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from . import cache_affinity # noqa: F401
|
||||
from . import cost_first # noqa: F401
|
||||
from . import free_first # noqa: F401
|
||||
from . import free_team_first # noqa: F401
|
||||
from . import health_first # noqa: F401
|
||||
from . import latency_first # noqa: F401
|
||||
from . import load_balance # noqa: F401
|
||||
from . import plus_first # noqa: F401
|
||||
from . import priority_first # noqa: F401
|
||||
from . import quota_balanced # noqa: F401
|
||||
from . import recent_refresh # noqa: F401
|
||||
from . import single_account # noqa: F401
|
||||
from . import team_first # noqa: F401
|
||||
from .registry import (
|
||||
PresetDimensionBase,
|
||||
PresetDimensionMeta,
|
||||
get_all_preset_dimensions,
|
||||
get_preset_dimension,
|
||||
get_preset_dimension_metas,
|
||||
get_preset_names,
|
||||
register_preset_dimension,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"PresetDimensionBase",
|
||||
"PresetDimensionMeta",
|
||||
"get_all_preset_dimensions",
|
||||
"get_preset_dimension",
|
||||
"get_preset_dimension_metas",
|
||||
"get_preset_names",
|
||||
"register_preset_dimension",
|
||||
]
|
||||
284
_deprecated_py_src/services/provider/pool/dimensions/_helpers.py
Normal file
284
_deprecated_py_src/services/provider/pool/dimensions/_helpers.py
Normal file
@@ -0,0 +1,284 @@
|
||||
"""Shared helpers for pool preset dimensions."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from src.core.provider_types import ProviderType, normalize_provider_type
|
||||
from src.services.provider_keys.quota_reader import get_quota_reader
|
||||
|
||||
|
||||
def safe_float(value: Any) -> float | None:
|
||||
try:
|
||||
parsed = float(value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
if math.isnan(parsed) or math.isinf(parsed):
|
||||
return None
|
||||
return parsed
|
||||
|
||||
|
||||
def safe_metadata(key_obj: Any) -> dict[str, Any]:
|
||||
raw = getattr(key_obj, "upstream_metadata", None)
|
||||
return raw if isinstance(raw, dict) else {}
|
||||
|
||||
|
||||
def normalize_plan(value: Any) -> str | None:
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
normalized = value.strip().lower()
|
||||
return normalized or None
|
||||
|
||||
|
||||
def rank_ascending(key_id: str, scores: dict[str, float], all_ids: list[str]) -> float:
|
||||
"""Rank score within all IDs; lower value means better rank."""
|
||||
|
||||
if not all_ids:
|
||||
return 0.0
|
||||
|
||||
valid_count = sum(1 for kid in all_ids if safe_float(scores.get(kid)) is not None)
|
||||
if valid_count <= 0:
|
||||
return 0.5
|
||||
|
||||
decorated: list[tuple[int, float, int, str]] = []
|
||||
for idx, kid in enumerate(all_ids):
|
||||
score_raw = safe_float(scores.get(kid))
|
||||
if score_raw is None:
|
||||
decorated.append((1, float("inf"), idx, kid))
|
||||
else:
|
||||
decorated.append((0, score_raw, idx, kid))
|
||||
|
||||
decorated.sort(key=lambda item: (item[0], item[1], item[2]))
|
||||
rank_idx = 0
|
||||
for idx, (_missing, _value, _order, kid) in enumerate(decorated):
|
||||
if kid == key_id:
|
||||
rank_idx = idx
|
||||
break
|
||||
|
||||
n = len(all_ids)
|
||||
if n <= 1:
|
||||
return 0.0
|
||||
return rank_idx / float(n - 1)
|
||||
|
||||
|
||||
def rank_descending(key_id: str, scores: dict[str, float], all_ids: list[str]) -> float:
|
||||
"""Rank score within all IDs; higher value means better rank."""
|
||||
|
||||
if not all_ids:
|
||||
return 0.0
|
||||
|
||||
valid_count = sum(1 for kid in all_ids if safe_float(scores.get(kid)) is not None)
|
||||
if valid_count <= 0:
|
||||
return 0.5
|
||||
|
||||
decorated: list[tuple[int, float, int, str]] = []
|
||||
for idx, kid in enumerate(all_ids):
|
||||
score_raw = safe_float(scores.get(kid))
|
||||
if score_raw is None:
|
||||
decorated.append((1, float("inf"), idx, kid))
|
||||
else:
|
||||
# 排序时取负值,使分值越大排名越靠前(rank 越小)
|
||||
decorated.append((0, -score_raw, idx, kid))
|
||||
|
||||
decorated.sort(key=lambda item: (item[0], item[1], item[2]))
|
||||
rank_idx = 0
|
||||
for idx, (_missing, _value, _order, kid) in enumerate(decorated):
|
||||
if kid == key_id:
|
||||
rank_idx = idx
|
||||
break
|
||||
|
||||
n = len(all_ids)
|
||||
if n <= 1:
|
||||
return 0.0
|
||||
return rank_idx / float(n - 1)
|
||||
|
||||
|
||||
def extract_plan_type(key_obj: Any) -> str | None:
|
||||
direct = normalize_plan(getattr(key_obj, "oauth_plan_type", None))
|
||||
if direct:
|
||||
return direct
|
||||
|
||||
metadata = safe_metadata(key_obj)
|
||||
for provider_type in (ProviderType.CODEX, ProviderType.KIRO, ProviderType.ANTIGRAVITY):
|
||||
plan_type = get_quota_reader(provider_type, metadata).plan_type()
|
||||
if plan_type:
|
||||
return plan_type
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_key_provider_type(key_obj: Any, provider_type: str | None = None) -> str | None:
|
||||
explicit = normalize_provider_type(provider_type)
|
||||
if explicit:
|
||||
return explicit
|
||||
|
||||
direct = normalize_provider_type(getattr(key_obj, "provider_type", None))
|
||||
if direct:
|
||||
return direct
|
||||
|
||||
provider = getattr(key_obj, "provider", None)
|
||||
related = normalize_provider_type(getattr(provider, "provider_type", None))
|
||||
if related:
|
||||
return related
|
||||
|
||||
metadata = safe_metadata(key_obj)
|
||||
candidates = [
|
||||
provider.value
|
||||
for provider in (ProviderType.CODEX, ProviderType.KIRO, ProviderType.ANTIGRAVITY)
|
||||
if isinstance(metadata.get(provider.value), dict)
|
||||
]
|
||||
if len(candidates) == 1:
|
||||
return candidates[0]
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _extract_codex_weekly_reset_seconds(metadata: dict[str, Any]) -> float | None:
|
||||
codex = metadata.get(ProviderType.CODEX.value)
|
||||
if not isinstance(codex, dict):
|
||||
return None
|
||||
|
||||
weekly_used_percent = safe_float(codex.get("primary_used_percent"))
|
||||
if weekly_used_percent is not None:
|
||||
clamped_used = max(0.0, min(weekly_used_percent, 100.0))
|
||||
if clamped_used <= 1e-6:
|
||||
# 周额度仍为满额时,不启用周窗口重置倒计时。
|
||||
return None
|
||||
|
||||
now = time.time()
|
||||
|
||||
# 优先绝对时间戳,避免 reset_seconds 快照随时间漂移。
|
||||
reset_at = safe_float(codex.get("primary_reset_at"))
|
||||
if reset_at is not None and reset_at > 0:
|
||||
remaining = reset_at - now
|
||||
return remaining if remaining > 0 else 0.0
|
||||
|
||||
reset_seconds = safe_float(codex.get("primary_reset_seconds"))
|
||||
if reset_seconds is None or reset_seconds < 0:
|
||||
return None
|
||||
|
||||
updated_at = safe_float(codex.get("updated_at"))
|
||||
if updated_at is not None and updated_at > 0:
|
||||
# 时钟偏移下 updated_at 可能晚于当前时间,elapsed 需要下限钳制到 0。
|
||||
elapsed = max(now - updated_at, 0.0)
|
||||
return max(reset_seconds - elapsed, 0.0)
|
||||
|
||||
return reset_seconds
|
||||
|
||||
|
||||
def extract_reset_seconds(key_obj: Any, provider_type: str | None = None) -> float | None:
|
||||
metadata = safe_metadata(key_obj)
|
||||
resolved_provider_type = _resolve_key_provider_type(key_obj, provider_type)
|
||||
|
||||
if resolved_provider_type == ProviderType.CODEX:
|
||||
# Codex metadata 已统一约定:primary_* 表示周限额,secondary_* 表示 5H 限额。
|
||||
return _extract_codex_weekly_reset_seconds(metadata)
|
||||
|
||||
if resolved_provider_type in (ProviderType.KIRO, ProviderType.ANTIGRAVITY):
|
||||
return get_quota_reader(resolved_provider_type, metadata).reset_seconds()
|
||||
|
||||
candidates: list[float] = []
|
||||
|
||||
for provider_type in (ProviderType.CODEX, ProviderType.KIRO, ProviderType.ANTIGRAVITY):
|
||||
reset_seconds = get_quota_reader(provider_type, metadata).reset_seconds()
|
||||
if reset_seconds is None:
|
||||
continue
|
||||
candidates.append(reset_seconds)
|
||||
|
||||
if not candidates:
|
||||
return None
|
||||
return min(candidates)
|
||||
|
||||
|
||||
def extract_usage_ratio(key_obj: Any) -> float | None:
|
||||
metadata = safe_metadata(key_obj)
|
||||
|
||||
for provider_type in (ProviderType.CODEX, ProviderType.KIRO, ProviderType.ANTIGRAVITY):
|
||||
usage_ratio = get_quota_reader(provider_type, metadata).usage_ratio()
|
||||
if usage_ratio is not None:
|
||||
return usage_ratio
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def extract_internal_priority(key_obj: Any) -> int:
|
||||
raw = getattr(key_obj, "internal_priority", None)
|
||||
parsed = safe_float(raw)
|
||||
if parsed is None:
|
||||
return 999999
|
||||
return max(0, int(parsed))
|
||||
|
||||
|
||||
def extract_health_score(key_obj: Any) -> float | None:
|
||||
direct = safe_float(getattr(key_obj, "health_score", None))
|
||||
if direct is not None:
|
||||
return max(0.0, min(direct, 1.0))
|
||||
|
||||
health_by_format = getattr(key_obj, "health_by_format", None)
|
||||
if not isinstance(health_by_format, dict) or not health_by_format:
|
||||
return None
|
||||
|
||||
scores: list[float] = []
|
||||
for payload in health_by_format.values():
|
||||
if not isinstance(payload, dict):
|
||||
continue
|
||||
score = safe_float(payload.get("health_score"))
|
||||
if score is None:
|
||||
continue
|
||||
scores.append(max(0.0, min(score, 1.0)))
|
||||
|
||||
if not scores:
|
||||
return None
|
||||
return min(scores)
|
||||
|
||||
|
||||
def plan_priority_score(plan_type: str | None, mode: str | None = None) -> float:
|
||||
"""Score a key based on plan type and scheduling mode.
|
||||
|
||||
Lower score = higher priority.
|
||||
"""
|
||||
|
||||
effective_mode = (mode or "both").strip().lower()
|
||||
if effective_mode == "free_only":
|
||||
if plan_type == "free":
|
||||
return 0.0
|
||||
if plan_type == "team":
|
||||
return 0.5
|
||||
elif effective_mode == "team_only":
|
||||
if plan_type == "team":
|
||||
return 0.0
|
||||
if plan_type == "free":
|
||||
return 0.5
|
||||
elif effective_mode == "plus_only":
|
||||
if plan_type in {"plus", "pro"}:
|
||||
return 0.0
|
||||
if plan_type in {"enterprise", "business"}:
|
||||
return 0.3
|
||||
else:
|
||||
# "both" or unrecognized -> original behavior
|
||||
if plan_type in {"free", "team"}:
|
||||
return 0.0
|
||||
if plan_type in {"enterprise", "business"}:
|
||||
return 0.2
|
||||
if plan_type in {"plus", "pro"}:
|
||||
return 0.6
|
||||
if plan_type:
|
||||
return 0.7
|
||||
return 0.8
|
||||
|
||||
|
||||
__all__ = [
|
||||
"extract_health_score",
|
||||
"extract_internal_priority",
|
||||
"extract_plan_type",
|
||||
"extract_reset_seconds",
|
||||
"extract_usage_ratio",
|
||||
"normalize_plan",
|
||||
"plan_priority_score",
|
||||
"rank_ascending",
|
||||
"rank_descending",
|
||||
"safe_float",
|
||||
"safe_metadata",
|
||||
]
|
||||
@@ -0,0 +1,45 @@
|
||||
"""cache_affinity preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import rank_descending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class CacheAffinityDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "cache_affinity"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "缓存亲和"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先复用最近使用过的 Key,利用 Prompt Caching"
|
||||
|
||||
@property
|
||||
def mutex_group(self) -> str | None:
|
||||
return "distribution_mode"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据 LRU 时间戳(最近使用优先,与 LRU 轮转相反)"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
return rank_descending(key_id, lru_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(CacheAffinityDimension())
|
||||
@@ -0,0 +1,62 @@
|
||||
"""cost_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_usage_ratio, rank_ascending, safe_float
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class CostFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "cost_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "成本优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先选择窗口消耗更低的账号"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据窗口成本/Token 用量,缺失时回退配额使用率"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
cost_totals = context.get("cost_totals")
|
||||
if not isinstance(cost_totals, dict):
|
||||
cost_totals = {}
|
||||
cost_limit = safe_float(context.get("cost_limit_per_key_tokens"))
|
||||
|
||||
cost_scores: dict[str, float] = {}
|
||||
for kid in all_key_ids:
|
||||
used = safe_float(cost_totals.get(kid))
|
||||
if used is not None and used >= 0:
|
||||
if cost_limit is not None and cost_limit > 0:
|
||||
cost_scores[kid] = max(0.0, min(used / cost_limit, 1.0))
|
||||
else:
|
||||
cost_scores[kid] = min(1.0, used / (used + 10000.0))
|
||||
continue
|
||||
|
||||
usage_ratio = extract_usage_ratio(keys_by_id.get(kid))
|
||||
if usage_ratio is not None:
|
||||
cost_scores[kid] = usage_ratio
|
||||
|
||||
if not cost_scores:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, cost_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(CostFirstDimension())
|
||||
@@ -0,0 +1,51 @@
|
||||
"""free_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_plan_type, plan_priority_score, rank_ascending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class FreeFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "free_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "Free 优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先消耗 Free 账号(依赖 plan_type)"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据 plan_type(Free 账号优先调度)"
|
||||
|
||||
@property
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
return ("codex", "kiro")
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
plan_scores = {
|
||||
kid: plan_priority_score(extract_plan_type(keys_by_id.get(kid)), "free_only")
|
||||
for kid in all_key_ids
|
||||
}
|
||||
if len(set(plan_scores.values())) <= 1:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, plan_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(FreeFirstDimension())
|
||||
@@ -0,0 +1,63 @@
|
||||
"""free_team_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_plan_type, plan_priority_score, rank_ascending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class FreeTeamFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "free_team_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "Free/Team 优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先消耗低档账号(依赖 plan_type)"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据 plan_type(oauth_plan_type 或 upstream_metadata)"
|
||||
|
||||
@property
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
return ("codex", "kiro")
|
||||
|
||||
@property
|
||||
def modes(self) -> tuple[str, ...] | None:
|
||||
return ("free_only", "team_only", "both")
|
||||
|
||||
@property
|
||||
def default_mode(self) -> str | None:
|
||||
return "both"
|
||||
|
||||
@property
|
||||
def hidden(self) -> bool:
|
||||
return True
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
plan_scores = {
|
||||
kid: plan_priority_score(extract_plan_type(keys_by_id.get(kid)), mode)
|
||||
for kid in all_key_ids
|
||||
}
|
||||
if len(set(plan_scores.values())) <= 1:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, plan_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(FreeTeamFirstDimension())
|
||||
@@ -0,0 +1,57 @@
|
||||
"""health_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_health_score, rank_ascending, safe_float
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class HealthFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "health_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "健康优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先选择健康分更高、失败更少的账号"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据 health_by_format 聚合分(含熔断/失败衰减)"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
health_scores_ctx = context.get("health_scores")
|
||||
if not isinstance(health_scores_ctx, dict):
|
||||
health_scores_ctx = {}
|
||||
|
||||
penalty_scores: dict[str, float] = {}
|
||||
for kid in all_key_ids:
|
||||
score = safe_float(health_scores_ctx.get(kid))
|
||||
if score is None:
|
||||
score = extract_health_score(keys_by_id.get(kid))
|
||||
if score is None:
|
||||
continue
|
||||
normalized = max(0.0, min(score, 1.0))
|
||||
penalty_scores[kid] = 1.0 - normalized
|
||||
|
||||
if not penalty_scores:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, penalty_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(HealthFirstDimension())
|
||||
@@ -0,0 +1,54 @@
|
||||
"""latency_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import rank_ascending, safe_float
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class LatencyFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "latency_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "延迟优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先选择最近延迟更低的账号"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据号池延迟窗口均值(latency_window_seconds)"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
latency_avgs = context.get("latency_avgs")
|
||||
if not isinstance(latency_avgs, dict):
|
||||
latency_avgs = {}
|
||||
|
||||
latency_scores: dict[str, float] = {}
|
||||
for kid in all_key_ids:
|
||||
latency = safe_float(latency_avgs.get(kid))
|
||||
if latency is None or latency < 0:
|
||||
continue
|
||||
latency_scores[kid] = latency
|
||||
|
||||
if not latency_scores:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, latency_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(LatencyFirstDimension())
|
||||
@@ -0,0 +1,45 @@
|
||||
"""load_balance preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from typing import Any
|
||||
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class LoadBalanceDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "load_balance"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "负载均衡"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "随机分散 Key 使用,均匀分摊负载"
|
||||
|
||||
@property
|
||||
def mutex_group(self) -> str | None:
|
||||
return "distribution_mode"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "每次随机分值,实现完全均匀分散"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
return random.random()
|
||||
|
||||
|
||||
register_preset_dimension(LoadBalanceDimension())
|
||||
@@ -0,0 +1,51 @@
|
||||
"""plus_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_plan_type, plan_priority_score, rank_ascending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class PlusFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "plus_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "Plus 优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先消耗 Plus/Pro 账号(依赖 plan_type)"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据 plan_type(Plus/Pro 账号优先调度)"
|
||||
|
||||
@property
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
return ("codex", "kiro")
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
plan_scores = {
|
||||
kid: plan_priority_score(extract_plan_type(keys_by_id.get(kid)), "plus_only")
|
||||
for kid in all_key_ids
|
||||
}
|
||||
if len(set(plan_scores.values())) <= 1:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, plan_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(PlusFirstDimension())
|
||||
@@ -0,0 +1,46 @@
|
||||
"""priority_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_internal_priority, rank_ascending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class PriorityFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "priority_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "优先级优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "按账号优先级顺序调度(数字越小越优先)"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据 internal_priority(支持拖拽/手工编辑)"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
priority_scores = {
|
||||
kid: float(extract_internal_priority(keys_by_id.get(kid))) for kid in all_key_ids
|
||||
}
|
||||
if len(set(priority_scores.values())) <= 1:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, priority_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(PriorityFirstDimension())
|
||||
@@ -0,0 +1,61 @@
|
||||
"""quota_balanced preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_usage_ratio, rank_ascending, safe_float
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class QuotaBalancedDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "quota_balanced"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "额度平均"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先选额度消耗最少的账号"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据账号配额使用率;无配额时回退到窗口成本使用"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
usage_scores: dict[str, float] = {}
|
||||
cost_totals = context.get("cost_totals")
|
||||
if not isinstance(cost_totals, dict):
|
||||
cost_totals = {}
|
||||
cost_limit = safe_float(context.get("cost_limit_per_key_tokens"))
|
||||
for kid in all_key_ids:
|
||||
key_obj = keys_by_id.get(kid)
|
||||
usage_ratio = extract_usage_ratio(key_obj)
|
||||
if usage_ratio is None:
|
||||
used = safe_float(cost_totals.get(kid))
|
||||
if used is not None and used >= 0:
|
||||
if cost_limit is not None and cost_limit > 0:
|
||||
usage_ratio = max(0.0, min(used / cost_limit, 1.0))
|
||||
else:
|
||||
# 无明确上限时用 log 归一化,确保维度仍有区分能力。
|
||||
usage_ratio = min(1.0, used / (used + 10000.0))
|
||||
if usage_ratio is not None:
|
||||
usage_scores[kid] = usage_ratio
|
||||
if not usage_scores:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, usage_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(QuotaBalancedDimension())
|
||||
@@ -0,0 +1,53 @@
|
||||
"""recent_refresh preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_reset_seconds, rank_ascending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class RecentRefreshDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "recent_refresh"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "额度刷新优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先选即将刷新额度的账号"
|
||||
|
||||
@property
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
return ("codex", "kiro")
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据账号额度重置倒计时(next_reset / reset_seconds)"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
provider_type = context.get("provider_type")
|
||||
reset_scores: dict[str, float] = {}
|
||||
for kid in all_key_ids:
|
||||
reset_seconds = extract_reset_seconds(keys_by_id.get(kid), provider_type=provider_type)
|
||||
if reset_seconds is not None:
|
||||
reset_scores[kid] = reset_seconds
|
||||
if not reset_scores:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, reset_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(RecentRefreshDimension())
|
||||
265
_deprecated_py_src/services/provider/pool/dimensions/registry.py
Normal file
265
_deprecated_py_src/services/provider/pool/dimensions/registry.py
Normal file
@@ -0,0 +1,265 @@
|
||||
"""Preset dimension registry for pool multi-score scheduling."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from threading import RLock
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PresetDimensionMeta:
|
||||
"""Serializable metadata for one preset dimension."""
|
||||
|
||||
name: str
|
||||
label: str
|
||||
description: str
|
||||
providers: tuple[str, ...]
|
||||
modes: tuple[str, ...] | None
|
||||
default_mode: str | None
|
||||
mutex_group: str | None
|
||||
evidence_hint: str | None
|
||||
|
||||
|
||||
class PresetDimensionBase(ABC):
|
||||
"""Base class of one pool scheduling preset dimension."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def name(self) -> str:
|
||||
"""Stable preset key, e.g. ``free_team_first``."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def label(self) -> str:
|
||||
"""User-facing label."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def description(self) -> str:
|
||||
"""User-facing description."""
|
||||
|
||||
@property
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
"""Supported provider types.
|
||||
|
||||
Empty tuple means the dimension is universal and applies to all providers.
|
||||
"""
|
||||
|
||||
return ()
|
||||
|
||||
@property
|
||||
def modes(self) -> tuple[str, ...] | None:
|
||||
"""Optional sub-modes for this dimension."""
|
||||
|
||||
return None
|
||||
|
||||
@property
|
||||
def default_mode(self) -> str | None:
|
||||
"""Default mode when mode is omitted."""
|
||||
|
||||
return None
|
||||
|
||||
@property
|
||||
def mutex_group(self) -> str | None:
|
||||
"""Optional mutual-exclusion group key.
|
||||
|
||||
Presets in the same group are expected to be mutually exclusive in UI.
|
||||
"""
|
||||
|
||||
return None
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
"""Human-readable hint about which data this preset uses."""
|
||||
|
||||
return None
|
||||
|
||||
@property
|
||||
def hidden(self) -> bool:
|
||||
"""If True, this dimension is excluded from API metadata listings.
|
||||
|
||||
The dimension remains functional for backward compatibility but
|
||||
will not appear in the scheduling dialog.
|
||||
"""
|
||||
|
||||
return False
|
||||
|
||||
@abstractmethod
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
"""Compute normalized metric in [0, 1], lower is better."""
|
||||
|
||||
def is_applicable(self, provider_type: str) -> bool:
|
||||
"""Return whether this dimension applies to the given provider type."""
|
||||
|
||||
if not self.providers:
|
||||
return True
|
||||
normalized = _normalize_name(provider_type)
|
||||
return normalized in self.providers
|
||||
|
||||
|
||||
def _normalize_name(value: Any) -> str:
|
||||
if not isinstance(value, str):
|
||||
return ""
|
||||
return value.strip().lower()
|
||||
|
||||
|
||||
def _normalize_names(values: tuple[str, ...] | list[str]) -> tuple[str, ...]:
|
||||
normalized = [_normalize_name(item) for item in values]
|
||||
return tuple(item for item in normalized if item)
|
||||
|
||||
|
||||
_registry_lock = RLock()
|
||||
_registry: dict[str, PresetDimensionBase] = {}
|
||||
|
||||
|
||||
def register_preset_dimension(dim: PresetDimensionBase) -> None:
|
||||
"""Register or replace one preset dimension by name."""
|
||||
|
||||
name = _normalize_name(dim.name)
|
||||
if not name:
|
||||
raise ValueError("preset dimension name must be a non-empty string")
|
||||
|
||||
providers = _normalize_names(dim.providers)
|
||||
modes = _normalize_names(dim.modes or ())
|
||||
default_mode = _normalize_name(dim.default_mode)
|
||||
|
||||
if modes and default_mode and default_mode not in modes:
|
||||
raise ValueError(f"default_mode must be one of modes for preset '{name}'")
|
||||
|
||||
class _NormalizedDimension(PresetDimensionBase):
|
||||
# Lightweight wrapper to keep normalized metadata while preserving compute logic.
|
||||
def __init__(self, wrapped: PresetDimensionBase) -> None:
|
||||
self._wrapped = wrapped
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return name
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return self._wrapped.label
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return self._wrapped.description
|
||||
|
||||
@property
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
return providers
|
||||
|
||||
@property
|
||||
def modes(self) -> tuple[str, ...] | None:
|
||||
return modes or None
|
||||
|
||||
@property
|
||||
def default_mode(self) -> str | None:
|
||||
if not modes:
|
||||
return None
|
||||
if default_mode:
|
||||
return default_mode
|
||||
return modes[0]
|
||||
|
||||
@property
|
||||
def mutex_group(self) -> str | None:
|
||||
raw = _normalize_name(self._wrapped.mutex_group)
|
||||
return raw or None
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
raw = str(self._wrapped.evidence_hint or "").strip()
|
||||
return raw or None
|
||||
|
||||
@property
|
||||
def hidden(self) -> bool:
|
||||
return self._wrapped.hidden
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
return self._wrapped.compute_metric(
|
||||
key_id=key_id,
|
||||
all_key_ids=all_key_ids,
|
||||
keys_by_id=keys_by_id,
|
||||
lru_scores=lru_scores,
|
||||
context=context,
|
||||
mode=mode,
|
||||
)
|
||||
|
||||
normalized = _NormalizedDimension(dim)
|
||||
with _registry_lock:
|
||||
_registry[name] = normalized
|
||||
|
||||
|
||||
def get_preset_dimension(name: str) -> PresetDimensionBase | None:
|
||||
"""Get one registered preset dimension by name."""
|
||||
|
||||
key = _normalize_name(name)
|
||||
if not key:
|
||||
return None
|
||||
with _registry_lock:
|
||||
return _registry.get(key)
|
||||
|
||||
|
||||
def get_all_preset_dimensions() -> list[PresetDimensionBase]:
|
||||
"""Get all registered preset dimensions in registration order."""
|
||||
|
||||
with _registry_lock:
|
||||
return list(_registry.values())
|
||||
|
||||
|
||||
def get_preset_names() -> set[str]:
|
||||
"""Get all registered preset names."""
|
||||
|
||||
with _registry_lock:
|
||||
return set(_registry.keys())
|
||||
|
||||
|
||||
def get_preset_dimension_metas() -> list[PresetDimensionMeta]:
|
||||
"""Get serializable metadata for all preset dimensions."""
|
||||
|
||||
metas: list[PresetDimensionMeta] = []
|
||||
for dim in get_all_preset_dimensions():
|
||||
if dim.hidden:
|
||||
continue
|
||||
metas.append(
|
||||
PresetDimensionMeta(
|
||||
name=dim.name,
|
||||
label=dim.label,
|
||||
description=dim.description,
|
||||
providers=dim.providers,
|
||||
modes=dim.modes,
|
||||
default_mode=dim.default_mode,
|
||||
mutex_group=dim.mutex_group,
|
||||
evidence_hint=dim.evidence_hint,
|
||||
)
|
||||
)
|
||||
return metas
|
||||
|
||||
|
||||
__all__ = [
|
||||
"PresetDimensionBase",
|
||||
"PresetDimensionMeta",
|
||||
"get_all_preset_dimensions",
|
||||
"get_preset_dimension",
|
||||
"get_preset_dimension_metas",
|
||||
"get_preset_names",
|
||||
"register_preset_dimension",
|
||||
]
|
||||
@@ -0,0 +1,51 @@
|
||||
"""single_account preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_internal_priority, rank_ascending, rank_descending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class SingleAccountDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "single_account"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "单号优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "集中使用同一账号(反向 LRU)"
|
||||
|
||||
@property
|
||||
def mutex_group(self) -> str | None:
|
||||
return "distribution_mode"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "先按账号优先级(internal_priority),同级再按反向 LRU 集中"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
priority_scores = {
|
||||
kid: float(extract_internal_priority(keys_by_id.get(kid))) for kid in all_key_ids
|
||||
}
|
||||
priority_rank = rank_ascending(key_id, priority_scores, all_key_ids)
|
||||
lru_concentrate_rank = rank_descending(key_id, lru_scores, all_key_ids)
|
||||
# 强化“单号优先”的可控性:优先级优先,反向 LRU 作为次级聚合。
|
||||
return max(0.0, min(priority_rank * 0.75 + lru_concentrate_rank * 0.25, 1.0))
|
||||
|
||||
|
||||
register_preset_dimension(SingleAccountDimension())
|
||||
@@ -0,0 +1,51 @@
|
||||
"""team_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_plan_type, plan_priority_score, rank_ascending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class TeamFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "team_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "Team 优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先消耗 Team 账号(依赖 plan_type)"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据 plan_type(Team 账号优先调度)"
|
||||
|
||||
@property
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
return ("codex", "kiro")
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
plan_scores = {
|
||||
kid: plan_priority_score(extract_plan_type(keys_by_id.get(kid)), "team_only")
|
||||
for kid in all_key_ids
|
||||
}
|
||||
if len(set(plan_scores.values())) <= 1:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, plan_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(TeamFirstDimension())
|
||||
104
_deprecated_py_src/services/provider/pool/health_cache.py
Normal file
104
_deprecated_py_src/services/provider/pool/health_cache.py
Normal file
@@ -0,0 +1,104 @@
|
||||
"""In-process pool health score cache.
|
||||
|
||||
This cache avoids recomputing per-key health aggregation on every request.
|
||||
It does not replace persistent health storage; source data still comes from
|
||||
``ProviderAPIKey.health_by_format`` carried on key objects.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
_TTL_SECONDS = 30.0
|
||||
_MAX_PROVIDERS = 500 # 最大缓存 provider 数量,防止无界增长
|
||||
_LOCK = threading.Lock()
|
||||
_CACHE: dict[str, tuple[float, dict[str, float]]] = {}
|
||||
|
||||
|
||||
def aggregate_health_score(health_by_format: Any) -> float:
|
||||
"""Aggregate health score from ``health_by_format`` (lower-bound strategy)."""
|
||||
if not isinstance(health_by_format, dict) or not health_by_format:
|
||||
return 1.0
|
||||
scores: list[float] = []
|
||||
for item in health_by_format.values():
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
try:
|
||||
score = float(item.get("health_score") or 1.0)
|
||||
except (TypeError, ValueError):
|
||||
score = 1.0
|
||||
scores.append(max(0.0, min(score, 1.0)))
|
||||
if not scores:
|
||||
return 1.0
|
||||
return min(scores)
|
||||
|
||||
|
||||
def get_health_scores(provider_id: str, keys: list[Any]) -> dict[str, float]:
|
||||
"""Return key health scores with per-provider TTL cache.
|
||||
|
||||
Uses incremental merge: if the cache is still valid but missing some keys,
|
||||
only the missing keys are computed and merged into the existing cache entry.
|
||||
"""
|
||||
now = time.monotonic()
|
||||
keys_by_id: dict[str, Any] = {}
|
||||
for k in keys:
|
||||
kid = str(getattr(k, "id", "") or "")
|
||||
if kid:
|
||||
keys_by_id[kid] = k
|
||||
if not keys_by_id:
|
||||
return {}
|
||||
|
||||
with _LOCK:
|
||||
cached = _CACHE.get(provider_id)
|
||||
if cached is not None:
|
||||
expires_at, payload = cached
|
||||
if now < expires_at:
|
||||
missing_ids = [kid for kid in keys_by_id if kid not in payload]
|
||||
if not missing_ids:
|
||||
return {kid: payload[kid] for kid in keys_by_id}
|
||||
# Compute only for missing keys, merge into existing cache
|
||||
for kid in missing_ids:
|
||||
payload[kid] = aggregate_health_score(
|
||||
getattr(keys_by_id[kid], "health_by_format", None)
|
||||
)
|
||||
# Trim stale key entries that are no longer in the current key set
|
||||
stale_ids = [sid for sid in payload if sid not in keys_by_id]
|
||||
for sid in stale_ids:
|
||||
del payload[sid]
|
||||
return {kid: payload[kid] for kid in keys_by_id}
|
||||
|
||||
fresh: dict[str, float] = {}
|
||||
for kid, key in keys_by_id.items():
|
||||
fresh[kid] = aggregate_health_score(getattr(key, "health_by_format", None))
|
||||
|
||||
with _LOCK:
|
||||
_CACHE[provider_id] = (now + _TTL_SECONDS, fresh)
|
||||
# 超出上限时清理过期条目,仍超限则淘汰最旧条目
|
||||
if len(_CACHE) > _MAX_PROVIDERS:
|
||||
expired = [k for k, (exp, _) in _CACHE.items() if now >= exp]
|
||||
for k in expired:
|
||||
del _CACHE[k]
|
||||
if len(_CACHE) > _MAX_PROVIDERS:
|
||||
oldest_key = min(_CACHE, key=lambda k: _CACHE[k][0])
|
||||
del _CACHE[oldest_key]
|
||||
return dict(fresh)
|
||||
|
||||
|
||||
def invalidate_provider_health_scores(provider_id: str) -> None:
|
||||
"""Invalidate health-score cache for one provider."""
|
||||
with _LOCK:
|
||||
_CACHE.pop(provider_id, None)
|
||||
|
||||
|
||||
def _clear_cache_for_tests() -> None:
|
||||
with _LOCK:
|
||||
_CACHE.clear()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"aggregate_health_score",
|
||||
"get_health_scores",
|
||||
"invalidate_provider_health_scores",
|
||||
]
|
||||
359
_deprecated_py_src/services/provider/pool/health_policy.py
Normal file
359
_deprecated_py_src/services/provider/pool/health_policy.py
Normal file
@@ -0,0 +1,359 @@
|
||||
"""Account Pool health policy: error code classification and key state management.
|
||||
|
||||
Maps upstream HTTP status codes to pool-level actions:
|
||||
|
||||
| Code | Action |
|
||||
|--------------|------------------------------------------------------------|
|
||||
| 401 | Invalidate OAuth token cache; permanent (deactivated) 1h, else no cooldown |
|
||||
| 402 | Long cooldown (payment issue) |
|
||||
| 403 | Graded cooldown: severe (suspended/banned) 1h, else 300s+ |
|
||||
| 400 | Check body for "organization has been disabled" -> cooldown |
|
||||
| 429 | Cooldown (retry-after or rate_limit_cooldown_seconds) |
|
||||
| 529 | Cooldown (overload_cooldown_seconds) |
|
||||
| * | Check unschedulable_rules keyword matching |
|
||||
| 408/5xx/etc | Transient cooldown (overload_cooldown_seconds) |
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.pool import redis_ops
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
|
||||
# Patterns in 400 error body that indicate account-level issues.
|
||||
_ACCOUNT_DISABLE_PATTERNS = (
|
||||
"organization has been disabled",
|
||||
"organization_disabled",
|
||||
"account has been disabled",
|
||||
"account_disabled",
|
||||
"account has been deactivated",
|
||||
"account_deactivated",
|
||||
"account deactivated",
|
||||
)
|
||||
|
||||
# 需要更长冷却的账号异常语义(403 body 关键字)。
|
||||
_FORBIDDEN_ACCOUNT_PATTERNS = (
|
||||
"account suspended",
|
||||
"account banned",
|
||||
"account deactivated",
|
||||
"subscription inactive",
|
||||
"suspended",
|
||||
"banned",
|
||||
"deactivated",
|
||||
)
|
||||
|
||||
_TRANSIENT_STATUS_COOLDOWN_REASON: dict[int, str] = {
|
||||
408: "request_timeout_408",
|
||||
409: "conflict_409",
|
||||
423: "locked_423",
|
||||
425: "too_early_425",
|
||||
500: "server_error_500",
|
||||
502: "bad_gateway_502",
|
||||
503: "service_unavailable_503",
|
||||
504: "gateway_timeout_504",
|
||||
}
|
||||
|
||||
|
||||
def _parse_retry_after(headers: dict[str, str] | None) -> int | None:
|
||||
"""Extract retry-after seconds from response headers."""
|
||||
if not headers:
|
||||
return None
|
||||
raw = headers.get("retry-after") or headers.get("Retry-After")
|
||||
if not raw:
|
||||
return None
|
||||
try:
|
||||
val = int(raw)
|
||||
return max(1, min(val, 3600))
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _extract_error_message(error_body: str | None) -> str:
|
||||
"""Best-effort extraction of error message from JSON body."""
|
||||
if not error_body:
|
||||
return ""
|
||||
try:
|
||||
data = json.loads(error_body)
|
||||
if isinstance(data, dict):
|
||||
error_obj = data.get("error")
|
||||
if isinstance(error_obj, dict):
|
||||
return str(error_obj.get("message", ""))
|
||||
if isinstance(error_obj, str):
|
||||
return error_obj
|
||||
return str(data.get("message", ""))
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
return error_body[:500]
|
||||
|
||||
|
||||
def _resolve_transient_cooldown_ttl(
|
||||
*,
|
||||
status_code: int,
|
||||
retry_after_seconds: int | None,
|
||||
config: PoolConfig,
|
||||
) -> int:
|
||||
"""Resolve cooldown ttl for transient upstream status codes."""
|
||||
if status_code in (429, 503):
|
||||
if retry_after_seconds is not None:
|
||||
return retry_after_seconds
|
||||
if status_code == 429:
|
||||
return config.rate_limit_cooldown_seconds
|
||||
# 408/409/423/425/5xx: 统一走短时过载冷却,避免雪崩重试。
|
||||
return config.overload_cooldown_seconds
|
||||
|
||||
|
||||
def _parse_google_quota_cooldown(error_body: str | None) -> int | None:
|
||||
"""Parse Google-specific quota cooldown from error body.
|
||||
|
||||
Safe to call for any provider: returns None unless the error body
|
||||
contains Google-specific fields (quotaResetTimeStamp / quotaResetDelay /
|
||||
"reset after" message pattern).
|
||||
"""
|
||||
if not error_body:
|
||||
return None
|
||||
try:
|
||||
from src.services.provider.adapters.gemini_cli.quota import extract_quota_cooldown_seconds
|
||||
|
||||
return extract_quota_cooldown_seconds(error_body)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def apply_health_policy(
|
||||
*,
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
status_code: int,
|
||||
error_body: str | None,
|
||||
response_headers: dict[str, str] | None,
|
||||
config: PoolConfig,
|
||||
) -> None:
|
||||
"""Apply health policy for an upstream error.
|
||||
|
||||
This is fire-and-forget; exceptions are caught and logged.
|
||||
"""
|
||||
if not config.health_policy_enabled:
|
||||
return
|
||||
|
||||
try:
|
||||
await _apply(
|
||||
provider_id=provider_id,
|
||||
key_id=key_id,
|
||||
status_code=status_code,
|
||||
error_body=error_body,
|
||||
response_headers=response_headers,
|
||||
config=config,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Pool health policy failed for key {}: {}",
|
||||
key_id[:8],
|
||||
str(exc),
|
||||
)
|
||||
|
||||
|
||||
async def _apply(
|
||||
*,
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
status_code: int,
|
||||
error_body: str | None,
|
||||
response_headers: dict[str, str] | None,
|
||||
config: PoolConfig,
|
||||
) -> None:
|
||||
error_msg = _extract_error_message(error_body)
|
||||
|
||||
# --- 401 Unauthorized ---------------------------------------------------
|
||||
if status_code == 401:
|
||||
await redis_ops.invalidate_oauth_token_cache(key_id)
|
||||
# Check if the 401 body indicates a permanent account-level deactivation
|
||||
# (e.g. OpenAI "account_deactivated"). These deserve a long cooldown.
|
||||
error_lower = error_msg.lower()
|
||||
is_permanent = any(p in error_lower for p in _ACCOUNT_DISABLE_PATTERNS)
|
||||
if is_permanent:
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "account_deactivated_401", ttl=3600)
|
||||
logger.warning(
|
||||
"Pool[{}]: key {} got 401 with account deactivation, cooldown 1h",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
)
|
||||
else:
|
||||
# Transient auth failure (e.g. expired token) — token cache already
|
||||
# invalidated above; the next request will trigger a token refresh.
|
||||
# No cooldown needed: the key should be retried immediately after refresh.
|
||||
logger.info(
|
||||
"Pool[{}]: key {} got 401, token cache invalidated (no cooldown)",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
)
|
||||
return
|
||||
|
||||
# --- 402 Payment Required ------------------------------------------------
|
||||
if status_code == 402:
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "payment_required_402", ttl=3600)
|
||||
logger.warning(
|
||||
"Pool[{}]: key {} got 402 (payment required), cooldown 1h",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
)
|
||||
return
|
||||
|
||||
# --- 403 Forbidden -------------------------------------------------------
|
||||
if status_code == 403:
|
||||
error_lower = error_msg.lower()
|
||||
severe = any(pattern in error_lower for pattern in _FORBIDDEN_ACCOUNT_PATTERNS)
|
||||
ttl = 3600 if severe else max(config.rate_limit_cooldown_seconds, 300)
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "forbidden_403", ttl=ttl)
|
||||
logger.warning(
|
||||
"Pool[{}]: key {} got 403 (forbidden), cooldown {}s",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
ttl,
|
||||
)
|
||||
return
|
||||
|
||||
# --- 400 with account-disable pattern ------------------------------------
|
||||
if status_code == 400:
|
||||
error_lower = error_msg.lower()
|
||||
for pattern in _ACCOUNT_DISABLE_PATTERNS:
|
||||
if pattern in error_lower:
|
||||
await redis_ops.set_cooldown(
|
||||
provider_id, key_id, f"account_disabled_400:{pattern}", ttl=3600
|
||||
)
|
||||
logger.warning(
|
||||
"Pool[{}]: key {} got 400 with '{}', cooldown 1h",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
pattern,
|
||||
)
|
||||
return
|
||||
|
||||
# --- 429 Rate Limited ----------------------------------------------------
|
||||
if status_code == 429:
|
||||
retry_after = _parse_retry_after(response_headers)
|
||||
if retry_after is None:
|
||||
retry_after = _parse_google_quota_cooldown(error_body)
|
||||
ttl = _resolve_transient_cooldown_ttl(
|
||||
status_code=status_code,
|
||||
retry_after_seconds=retry_after,
|
||||
config=config,
|
||||
)
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "rate_limited_429", ttl=ttl)
|
||||
logger.info(
|
||||
"Pool[{}]: key {} got 429, cooldown {}s",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
ttl,
|
||||
)
|
||||
return
|
||||
|
||||
# --- 529 Overloaded ------------------------------------------------------
|
||||
if status_code == 529:
|
||||
ttl = config.overload_cooldown_seconds
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "overloaded_529", ttl=ttl)
|
||||
logger.info(
|
||||
"Pool[{}]: key {} got 529, cooldown {}s",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
ttl,
|
||||
)
|
||||
return
|
||||
|
||||
# --- Keyword-based unschedulable rules -----------------------------------
|
||||
if config.unschedulable_rules and error_msg:
|
||||
error_lower = error_msg.lower()
|
||||
for rule in config.unschedulable_rules:
|
||||
if rule.keyword.lower() in error_lower:
|
||||
ttl = max(60, rule.duration_minutes * 60)
|
||||
await redis_ops.set_cooldown(
|
||||
provider_id,
|
||||
key_id,
|
||||
f"rule:{rule.keyword}",
|
||||
ttl=ttl,
|
||||
)
|
||||
logger.info(
|
||||
"Pool[{}]: key {} matched rule '{}', cooldown {}m",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
rule.keyword,
|
||||
rule.duration_minutes,
|
||||
)
|
||||
return
|
||||
|
||||
# --- Transient status bucket (408/409/423/425/5xx) ----------------------
|
||||
reason = _TRANSIENT_STATUS_COOLDOWN_REASON.get(status_code)
|
||||
if reason:
|
||||
retry_after = _parse_retry_after(response_headers)
|
||||
ttl = _resolve_transient_cooldown_ttl(
|
||||
status_code=status_code,
|
||||
retry_after_seconds=retry_after,
|
||||
config=config,
|
||||
)
|
||||
await redis_ops.set_cooldown(provider_id, key_id, reason, ttl=ttl)
|
||||
logger.info(
|
||||
"Pool[{}]: key {} got {}, cooldown {}s ({})",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
status_code,
|
||||
ttl,
|
||||
reason,
|
||||
)
|
||||
return
|
||||
|
||||
|
||||
async def apply_stream_timeout_policy(
|
||||
*,
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
config: PoolConfig,
|
||||
) -> None:
|
||||
"""Record a stream timeout event and apply cooldown if threshold is reached.
|
||||
|
||||
Called when an upstream stream response times out (no data within the
|
||||
configured interval). Increments a per-key counter in Redis and sets
|
||||
a cooldown if the count reaches the configured threshold.
|
||||
"""
|
||||
if not config.health_policy_enabled:
|
||||
return
|
||||
|
||||
try:
|
||||
count = await redis_ops.incr_stream_timeout_count(
|
||||
provider_id,
|
||||
key_id,
|
||||
config.stream_timeout_window_seconds,
|
||||
)
|
||||
if count >= config.stream_timeout_threshold:
|
||||
ttl = config.stream_timeout_cooldown_seconds
|
||||
await redis_ops.set_cooldown(
|
||||
provider_id,
|
||||
key_id,
|
||||
f"stream_timeout_x{count}",
|
||||
ttl=ttl,
|
||||
)
|
||||
logger.warning(
|
||||
"Pool[{}]: key {} stream timeout count {} >= threshold {}, cooldown {}s",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
count,
|
||||
config.stream_timeout_threshold,
|
||||
ttl,
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"Pool[{}]: key {} stream timeout count {}/{}",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
count,
|
||||
config.stream_timeout_threshold,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Pool stream timeout policy failed for key {}: {}",
|
||||
key_id[:8],
|
||||
str(exc),
|
||||
)
|
||||
91
_deprecated_py_src/services/provider/pool/hooks.py
Normal file
91
_deprecated_py_src/services/provider/pool/hooks.py
Normal file
@@ -0,0 +1,91 @@
|
||||
"""Pool scheduling hooks -- provider-type-specific pool behaviour.
|
||||
|
||||
Some provider types need custom logic during pool scheduling (e.g. extracting
|
||||
a session UUID for sticky binding). This module provides a small Protocol +
|
||||
registry so the pool layer stays generic while provider-specific behaviour
|
||||
lives alongside each adapter.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from typing import Any, Protocol, runtime_checkable
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Protocol
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class PoolSchedulingHook(Protocol):
|
||||
"""Provider-type-specific pool scheduling behaviour.
|
||||
|
||||
Each provider type can optionally register a hook to customize:
|
||||
- Session UUID extraction (for sticky sessions)
|
||||
- Post-success / post-error callbacks
|
||||
|
||||
Optional methods (checked via ``hasattr`` by callers):
|
||||
- ``on_pool_success``
|
||||
- ``on_pool_error``
|
||||
"""
|
||||
|
||||
name: str
|
||||
|
||||
def extract_session_uuid(self, request_body: dict[str, Any]) -> str | None:
|
||||
"""Extract a session UUID for sticky binding from the request body."""
|
||||
...
|
||||
|
||||
# -- Optional lifecycle callbacks -----------------------------------------
|
||||
# These are checked via ``hasattr`` so existing implementations that
|
||||
# don't define them will continue to work.
|
||||
|
||||
def on_pool_success(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
session_uuid: str | None,
|
||||
context: dict[str, Any],
|
||||
) -> None:
|
||||
"""Called after a successful pool request (provider-specific logic)."""
|
||||
...
|
||||
|
||||
def on_pool_error(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
status_code: int,
|
||||
context: dict[str, Any],
|
||||
) -> None:
|
||||
"""Called after a failed pool request (provider-specific logic)."""
|
||||
...
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_hook_registry: dict[str, PoolSchedulingHook] = {}
|
||||
_registry_lock = threading.Lock()
|
||||
|
||||
|
||||
def register_pool_hook(provider_type: str, hook: PoolSchedulingHook) -> None:
|
||||
"""Register a pool scheduling hook for a provider type."""
|
||||
from src.core.provider_types import normalize_provider_type
|
||||
|
||||
pt = normalize_provider_type(provider_type)
|
||||
with _registry_lock:
|
||||
_hook_registry[pt] = hook
|
||||
|
||||
|
||||
def get_pool_hook(provider_type: str | None) -> PoolSchedulingHook | None:
|
||||
"""Return the pool scheduling hook for a provider type, or ``None``."""
|
||||
if not provider_type:
|
||||
return None
|
||||
from src.services.provider.envelope import ensure_providers_bootstrapped
|
||||
|
||||
ensure_providers_bootstrapped(provider_types=[provider_type])
|
||||
|
||||
from src.core.provider_types import normalize_provider_type
|
||||
|
||||
pt = normalize_provider_type(provider_type)
|
||||
return _hook_registry.get(pt)
|
||||
621
_deprecated_py_src/services/provider/pool/manager.py
Normal file
621
_deprecated_py_src/services/provider/pool/manager.py
Normal file
@@ -0,0 +1,621 @@
|
||||
"""Account Pool Manager (provider-agnostic).
|
||||
|
||||
Stateless facade that coordinates pool operations for any Provider with
|
||||
pool configuration enabled. All state lives in Redis via :mod:`redis_ops`.
|
||||
|
||||
Usage::
|
||||
|
||||
mgr = PoolManager(provider_id, pool_config)
|
||||
reordered = await mgr.reorder_candidates(session_uuid, candidates)
|
||||
# ... execute request ...
|
||||
await mgr.on_request_success(session_uuid=..., key_id=..., tokens_used=...)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import random
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any, TypeVar
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.pool import redis_ops
|
||||
from src.services.provider.pool.account_state import resolve_pool_account_state
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
from src.services.provider.pool.health_cache import get_health_scores
|
||||
from src.services.provider.pool.trace import PoolCandidateTrace, PoolSchedulingTrace
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.models.database import ProviderAPIKey
|
||||
from src.services.scheduling.schemas import ProviderCandidate
|
||||
|
||||
|
||||
class PoolManager:
|
||||
"""Coordinate pool-level scheduling for a single Provider."""
|
||||
|
||||
__slots__ = ("provider_id", "config", "provider_type")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider_id: str,
|
||||
config: PoolConfig,
|
||||
provider_type: str | None = None,
|
||||
) -> None:
|
||||
self.provider_id = provider_id
|
||||
self.config = config
|
||||
self.provider_type = str(provider_type or "").strip().lower() or None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Core scheduling: reorder candidate list for pool-aware selection
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def reorder_candidates(
|
||||
self,
|
||||
session_uuid: str | None,
|
||||
candidates: list[ProviderCandidate],
|
||||
) -> list[ProviderCandidate]:
|
||||
"""Reorder *candidates* according to pool rules.
|
||||
|
||||
The returned list keeps the same elements but in a new order:
|
||||
|
||||
1. **Sticky session hit** -- if the session is already bound to a key
|
||||
and that key appears in *candidates* and is not in cooldown, move it
|
||||
to position 0.
|
||||
2. **Filter** out keys in account-blocked / cooldown / cost-exhausted state (mark
|
||||
``is_skipped``).
|
||||
3. **LRU sort** -- among remaining candidates at the same priority
|
||||
level, sort by least-recently-used.
|
||||
4. **Random tiebreak** -- among candidates with identical LRU score.
|
||||
|
||||
Also builds a :class:`PoolSchedulingTrace` and attaches per-candidate
|
||||
trace data via ``_pool_extra_data`` / ``_pool_scheduling_trace``
|
||||
attributes on candidate objects.
|
||||
"""
|
||||
if not candidates:
|
||||
return candidates
|
||||
|
||||
pid = self.provider_id
|
||||
|
||||
# Build trace
|
||||
trace = PoolSchedulingTrace(
|
||||
provider_id=pid,
|
||||
total_keys=len(candidates),
|
||||
session_uuid=session_uuid[:8] if session_uuid else None,
|
||||
)
|
||||
|
||||
# --- Strategy: before_select ----------------------------------
|
||||
strategies = _get_active_strategies(self.config)
|
||||
key_ids = [str(c.key.id) for c in candidates]
|
||||
strategy_context: dict[str, Any] = {"session_uuid": session_uuid}
|
||||
for strategy in strategies:
|
||||
if hasattr(strategy, "on_before_select"):
|
||||
try:
|
||||
filtered = strategy.on_before_select(
|
||||
provider_id=pid,
|
||||
key_ids=key_ids,
|
||||
config=self.config,
|
||||
context=strategy_context,
|
||||
)
|
||||
if filtered is not None:
|
||||
key_ids = filtered
|
||||
except Exception:
|
||||
logger.opt(exception=True).debug(
|
||||
"Pool[{}]: strategy before_select failed", pid[:8]
|
||||
)
|
||||
|
||||
# --- 1. Sticky session ----------------------------------------
|
||||
sticky_key_id: str | None = None
|
||||
if session_uuid and self.config.sticky_session_ttl_seconds > 0:
|
||||
sticky_key_id = await redis_ops.get_sticky_binding(
|
||||
pid, session_uuid, self.config.sticky_session_ttl_seconds
|
||||
)
|
||||
|
||||
provider_type = self.provider_type
|
||||
if provider_type is None and candidates:
|
||||
first_provider = getattr(candidates[0], "provider", None)
|
||||
provider_type = str(getattr(first_provider, "provider_type", "") or "").strip().lower()
|
||||
if not provider_type:
|
||||
provider_type = None
|
||||
|
||||
# --- 2. Batch fetch pool state (parallel) ---------------------
|
||||
all_key_ids = [str(c.key.id) for c in candidates]
|
||||
|
||||
# Fire independent Redis queries concurrently.
|
||||
# Only fetch reason (no TTL) on the scheduling hot path -- TTL is only
|
||||
# used for trace display and costs an extra pipeline command per key.
|
||||
_cooldown_coro = redis_ops.batch_get_cooldowns(pid, all_key_ids, include_ttl=False)
|
||||
_cost_coro = (
|
||||
redis_ops.batch_get_cost_totals(pid, all_key_ids, self.config.cost_window_seconds)
|
||||
if (
|
||||
self.config.cost_limit_per_key_tokens is not None
|
||||
or self.config.scheduling_mode == "multi_score"
|
||||
)
|
||||
else None
|
||||
)
|
||||
# LRU scores are needed both for plain LRU sorting and for multi_score
|
||||
# dimensions (e.g. cache_affinity / single_account) that rely on
|
||||
# lru_scores data.
|
||||
_need_lru = self.config.lru_enabled or self.config.scheduling_mode == "multi_score"
|
||||
_lru_coro = redis_ops.get_lru_scores(pid, all_key_ids) if _need_lru else None
|
||||
_latency_coro = (
|
||||
redis_ops.batch_get_latency_avgs(pid, all_key_ids, self.config.latency_window_seconds)
|
||||
if self.config.scheduling_mode == "multi_score"
|
||||
else None
|
||||
)
|
||||
|
||||
# Gather all non-None coroutines in parallel.
|
||||
coros: list[Any] = [_cooldown_coro]
|
||||
_cost_idx = -1
|
||||
_lru_idx = -1
|
||||
_latency_idx = -1
|
||||
if _cost_coro is not None:
|
||||
_cost_idx = len(coros)
|
||||
coros.append(_cost_coro)
|
||||
if _lru_coro is not None:
|
||||
_lru_idx = len(coros)
|
||||
coros.append(_lru_coro)
|
||||
if _latency_coro is not None:
|
||||
_latency_idx = len(coros)
|
||||
coros.append(_latency_coro)
|
||||
|
||||
gathered = await asyncio.gather(*coros)
|
||||
|
||||
cooldowns: dict[str, str | None] = gathered[0]
|
||||
|
||||
# Cost check
|
||||
cost_exhausted: set[str] = set()
|
||||
cost_soft: set[str] = set()
|
||||
cost_totals: dict[str, int] = {}
|
||||
if _cost_idx >= 0:
|
||||
cost_totals = gathered[_cost_idx]
|
||||
limit = self.config.cost_limit_per_key_tokens
|
||||
if limit is not None:
|
||||
for kid, total in cost_totals.items():
|
||||
if total >= limit:
|
||||
cost_exhausted.add(kid)
|
||||
elif total >= limit * self.config.cost_soft_threshold_percent / 100:
|
||||
cost_soft.add(kid)
|
||||
|
||||
# LRU scores
|
||||
lru_scores: dict[str, float] = {}
|
||||
if _lru_idx >= 0:
|
||||
lru_scores = gathered[_lru_idx]
|
||||
|
||||
# Latency averages
|
||||
latency_avgs: dict[str, float] = {}
|
||||
if _latency_idx >= 0:
|
||||
latency_avgs = gathered[_latency_idx]
|
||||
|
||||
# Health scores (TTL cached, no Redis round-trip) -- only needed for multi_score
|
||||
health_scores: dict[str, float] = {}
|
||||
if self.config.scheduling_mode == "multi_score":
|
||||
health_scores = get_health_scores(pid, [c.key for c in candidates])
|
||||
|
||||
strategy_context.update(
|
||||
{
|
||||
"provider_type": provider_type,
|
||||
"all_key_ids": all_key_ids,
|
||||
"lru_scores": lru_scores,
|
||||
"cost_totals": cost_totals,
|
||||
"cost_limit_per_key_tokens": self.config.cost_limit_per_key_tokens,
|
||||
"latency_avgs": latency_avgs,
|
||||
"health_scores": health_scores,
|
||||
"keys_by_id": {str(c.key.id): c.key for c in candidates},
|
||||
}
|
||||
)
|
||||
|
||||
# --- Strategy: compute_score ----------------------------------
|
||||
custom_scores: dict[str, float] = {}
|
||||
for strategy in strategies:
|
||||
if hasattr(strategy, "compute_score"):
|
||||
for kid in all_key_ids:
|
||||
try:
|
||||
custom = strategy.compute_score(
|
||||
key_id=kid,
|
||||
config=self.config,
|
||||
context=strategy_context,
|
||||
)
|
||||
if custom is not None:
|
||||
custom_scores[kid] = float(custom)
|
||||
lru_scores[kid] = float(custom)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# --- 3. Classify candidates -----------------------------------
|
||||
# Use precomputed account states from CandidateBuilder when available
|
||||
# (upstream_metadata is deferred on pool keys to save memory).
|
||||
# Fall back to on-the-fly resolution for non-pool candidates.
|
||||
account_states: dict[str, Any] = {}
|
||||
for c in candidates:
|
||||
kid = str(c.key.id)
|
||||
if kid not in account_states:
|
||||
precomputed = getattr(c.key, "_pool_account_state", None)
|
||||
if precomputed is not None:
|
||||
account_states[kid] = precomputed
|
||||
else:
|
||||
account_states[kid] = resolve_pool_account_state(
|
||||
provider_type=provider_type,
|
||||
upstream_metadata=getattr(c.key, "upstream_metadata", None),
|
||||
oauth_invalid_reason=getattr(c.key, "oauth_invalid_reason", None),
|
||||
)
|
||||
|
||||
sticky_candidate: ProviderCandidate | None = None
|
||||
available: list[ProviderCandidate] = []
|
||||
skipped: list[ProviderCandidate] = []
|
||||
|
||||
for c in candidates:
|
||||
kid = str(c.key.id)
|
||||
ct = PoolCandidateTrace(key_id=kid)
|
||||
ct.scoring_mode = self.config.scheduling_mode
|
||||
ct.latency_avg_ms = float(latency_avgs.get(kid, 0.0) or 0.0)
|
||||
ct.health_score = float(health_scores.get(kid, 1.0) or 1.0)
|
||||
if kid in custom_scores:
|
||||
ct.composite_score = float(custom_scores[kid])
|
||||
|
||||
# Already skipped upstream?
|
||||
if c.is_skipped:
|
||||
skipped.append(c)
|
||||
ct.skipped = True
|
||||
ct.skip_type = "upstream"
|
||||
trace.candidate_traces[kid] = ct
|
||||
continue
|
||||
|
||||
# Account blocked?
|
||||
account_state = account_states[kid]
|
||||
if account_state.blocked:
|
||||
c.is_skipped = True
|
||||
skip_reason = account_state.reason or account_state.label or "account blocked"
|
||||
c.skip_reason = f"pool account blocked: {skip_reason}"
|
||||
skipped.append(c)
|
||||
ct.skipped = True
|
||||
ct.skip_type = "account_blocked"
|
||||
ct.account_block_code = account_state.code
|
||||
ct.account_block_label = account_state.label
|
||||
ct.account_block_reason = account_state.reason
|
||||
_attach_pool_extra(c, ct)
|
||||
trace.candidate_traces[kid] = ct
|
||||
continue
|
||||
|
||||
cd_reason = cooldowns.get(kid)
|
||||
if cd_reason is not None:
|
||||
c.is_skipped = True
|
||||
c.skip_reason = f"pool cooldown: {cd_reason}"
|
||||
skipped.append(c)
|
||||
ct.skipped = True
|
||||
ct.skip_type = "cooldown"
|
||||
ct.cooldown_reason = cd_reason
|
||||
ct.cooldown_ttl = None # TTL skipped on hot path for perf
|
||||
_attach_pool_extra(c, ct)
|
||||
trace.candidate_traces[kid] = ct
|
||||
continue
|
||||
|
||||
# Cost exhausted?
|
||||
if kid in cost_exhausted:
|
||||
c.is_skipped = True
|
||||
c.skip_reason = "pool cost limit reached"
|
||||
skipped.append(c)
|
||||
ct.skipped = True
|
||||
ct.skip_type = "cost_exhausted"
|
||||
ct.cost_window_usage = cost_totals.get(kid, 0)
|
||||
ct.cost_limit = self.config.cost_limit_per_key_tokens
|
||||
_attach_pool_extra(c, ct)
|
||||
trace.candidate_traces[kid] = ct
|
||||
continue
|
||||
|
||||
# Sticky hit?
|
||||
if sticky_key_id and kid == sticky_key_id:
|
||||
sticky_candidate = c
|
||||
ct.reason = "sticky"
|
||||
ct.sticky_hit = True
|
||||
trace.sticky_session_used = True
|
||||
else:
|
||||
available.append(c)
|
||||
if kid in custom_scores and self.config.scheduling_mode == "multi_score":
|
||||
ct.reason = "multi_score"
|
||||
else:
|
||||
ct.reason = "lru" if lru_scores.get(kid, 0) > 0 else "random"
|
||||
|
||||
ct.lru_score = lru_scores.get(kid, 0.0)
|
||||
ct.cost_window_usage = cost_totals.get(kid, 0)
|
||||
ct.cost_limit = self.config.cost_limit_per_key_tokens
|
||||
if kid in cost_soft:
|
||||
ct.cost_soft_threshold = True
|
||||
_attach_pool_extra(c, ct)
|
||||
trace.candidate_traces[kid] = ct
|
||||
|
||||
# --- 4. Sort available by LRU ---------------------------------
|
||||
if lru_scores and available:
|
||||
available.sort(key=lambda c: lru_scores.get(str(c.key.id), 0.0))
|
||||
|
||||
# Random tiebreak among candidates with the same LRU score
|
||||
if len(available) > 1 and lru_scores:
|
||||
_shuffle_same_score_groups(available, lru_scores)
|
||||
|
||||
# --- 5. Assemble final order ----------------------------------
|
||||
result: list[ProviderCandidate] = []
|
||||
if sticky_candidate is not None:
|
||||
result.append(sticky_candidate)
|
||||
result.extend(available)
|
||||
result.extend(skipped)
|
||||
|
||||
if sticky_candidate:
|
||||
logger.debug(
|
||||
"Pool[{}]: sticky hit key={}",
|
||||
pid[:8],
|
||||
sticky_key_id and sticky_key_id[:8],
|
||||
)
|
||||
|
||||
# --- Strategy: after_select -----------------------------------
|
||||
if result:
|
||||
first_kid = str(result[0].key.id)
|
||||
first_trace = trace.candidate_traces.get(first_kid)
|
||||
for strategy in strategies:
|
||||
if hasattr(strategy, "on_after_select") and first_trace:
|
||||
try:
|
||||
strategy.on_after_select(
|
||||
provider_id=pid,
|
||||
selected_key_id=first_kid,
|
||||
trace=first_trace,
|
||||
config=self.config,
|
||||
context=strategy_context,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Attach the full trace to the first candidate for downstream use.
|
||||
if result:
|
||||
setattr(result[0], "_pool_scheduling_trace", trace)
|
||||
|
||||
return result
|
||||
|
||||
async def select_pool_keys(
|
||||
self,
|
||||
session_uuid: str | None,
|
||||
keys: list[ProviderAPIKey],
|
||||
*,
|
||||
availability_checker: (
|
||||
Callable[[ProviderAPIKey], tuple[bool, str | None, str | None]] | None
|
||||
) = None,
|
||||
page_size: int = 50,
|
||||
) -> tuple[list[ProviderAPIKey], PoolSchedulingTrace]:
|
||||
"""Select and order pool keys with trace output.
|
||||
|
||||
Reuses :meth:`reorder_candidates` logic by adapting keys to lightweight
|
||||
candidate-like wrappers, then propagates skip/trace metadata back onto
|
||||
each key object for downstream execution/recording.
|
||||
|
||||
When *availability_checker* is provided, post-sort availability checks
|
||||
are performed lazily: only the top *page_size* non-skipped keys are
|
||||
checked at a time; if all fail, the next page is checked, and so on.
|
||||
Keys beyond the last checked page are marked as ``deferred`` (skipped
|
||||
without checking) to avoid unnecessary CPU work on large pools.
|
||||
"""
|
||||
if not keys:
|
||||
return (
|
||||
[],
|
||||
PoolSchedulingTrace(
|
||||
provider_id=self.provider_id,
|
||||
total_keys=0,
|
||||
session_uuid=session_uuid[:8] if session_uuid else None,
|
||||
),
|
||||
)
|
||||
|
||||
class _KeyCandidate:
|
||||
__slots__ = (
|
||||
"key",
|
||||
"is_skipped",
|
||||
"skip_reason",
|
||||
"_pool_extra_data",
|
||||
"_pool_scheduling_trace",
|
||||
)
|
||||
|
||||
def __init__(self, key: ProviderAPIKey) -> None:
|
||||
self.key = key
|
||||
self.is_skipped = False
|
||||
self.skip_reason: str | None = None
|
||||
self._pool_extra_data: dict | None = None
|
||||
self._pool_scheduling_trace: PoolSchedulingTrace | None = None
|
||||
|
||||
wrappers = [_KeyCandidate(k) for k in keys]
|
||||
reordered_wrappers = await self.reorder_candidates(session_uuid, wrappers) # type: ignore[arg-type]
|
||||
|
||||
trace: PoolSchedulingTrace | None = None
|
||||
if reordered_wrappers:
|
||||
maybe_trace = getattr(reordered_wrappers[0], "_pool_scheduling_trace", None)
|
||||
if isinstance(maybe_trace, PoolSchedulingTrace):
|
||||
trace = maybe_trace
|
||||
if trace is None:
|
||||
trace = PoolSchedulingTrace(
|
||||
provider_id=self.provider_id,
|
||||
total_keys=len(keys),
|
||||
session_uuid=session_uuid[:8] if session_uuid else None,
|
||||
)
|
||||
|
||||
ordered_keys: list[ProviderAPIKey] = []
|
||||
for order_idx, wrapped in enumerate(reordered_wrappers):
|
||||
key = wrapped.key
|
||||
is_skipped = bool(getattr(wrapped, "is_skipped", False))
|
||||
skip_reason = str(getattr(wrapped, "skip_reason", "") or "")
|
||||
setattr(key, "_pool_skipped", is_skipped)
|
||||
setattr(key, "_pool_skip_reason", skip_reason if skip_reason else None)
|
||||
setattr(key, "_pool_order_index", order_idx)
|
||||
pool_extra = getattr(wrapped, "_pool_extra_data", None)
|
||||
setattr(
|
||||
key, "_pool_extra_data", dict(pool_extra) if isinstance(pool_extra, dict) else {}
|
||||
)
|
||||
ordered_keys.append(key)
|
||||
|
||||
# -- 分页可用性检查 --
|
||||
# 排序后对非 skipped key 分页调用 availability_checker,
|
||||
# 找到 page_size 个可用 key 后停止检查,剩余标记 deferred。
|
||||
if availability_checker is not None:
|
||||
available_count = 0
|
||||
found_enough = False
|
||||
for key in ordered_keys:
|
||||
if getattr(key, "_pool_skipped", False):
|
||||
continue
|
||||
if found_enough:
|
||||
setattr(key, "_pool_skipped", True)
|
||||
setattr(key, "_pool_skip_reason", "deferred")
|
||||
continue
|
||||
is_available, skip_reason_check, mapping_model = availability_checker(key)
|
||||
if not is_available:
|
||||
setattr(key, "_pool_skipped", True)
|
||||
setattr(key, "_pool_skip_reason", skip_reason_check)
|
||||
else:
|
||||
if mapping_model:
|
||||
setattr(key, "_pool_mapping_matched_model", mapping_model)
|
||||
available_count += 1
|
||||
if available_count >= page_size:
|
||||
found_enough = True
|
||||
|
||||
return ordered_keys, trace
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Post-request hooks
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def on_request_success(
|
||||
self,
|
||||
*,
|
||||
session_uuid: str | None,
|
||||
key_id: str,
|
||||
tokens_used: int = 0,
|
||||
ttfb_ms: int | None = None,
|
||||
) -> None:
|
||||
"""Called after a successful upstream request."""
|
||||
pid = self.provider_id
|
||||
|
||||
# Bind sticky session
|
||||
if session_uuid and self.config.sticky_session_ttl_seconds > 0:
|
||||
await redis_ops.set_sticky_binding(
|
||||
pid, session_uuid, key_id, self.config.sticky_session_ttl_seconds
|
||||
)
|
||||
|
||||
# Touch LRU -- needed for both plain LRU mode and multi_score dimensions
|
||||
# (e.g. cache_affinity) that rely on LRU timestamps.
|
||||
if self.config.lru_enabled or self.config.scheduling_mode == "multi_score":
|
||||
await redis_ops.touch_lru(pid, key_id)
|
||||
|
||||
# Record cost
|
||||
if tokens_used > 0 and self.config.cost_limit_per_key_tokens is not None:
|
||||
await redis_ops.add_cost_entry(
|
||||
pid, key_id, tokens_used, self.config.cost_window_seconds
|
||||
)
|
||||
|
||||
# Record latency sample for multi-score scheduling.
|
||||
if self.config.scheduling_mode == "multi_score" and ttfb_ms is not None and ttfb_ms >= 0:
|
||||
await redis_ops.record_latency(
|
||||
pid,
|
||||
key_id,
|
||||
ttfb_ms,
|
||||
self.config.latency_window_seconds,
|
||||
self.config.latency_sample_limit,
|
||||
)
|
||||
|
||||
async def on_request_error(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
status_code: int,
|
||||
error_body: str | None = None,
|
||||
response_headers: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
"""Called after an upstream error. Delegates to health policy."""
|
||||
# Import lazily to avoid circular deps
|
||||
from src.services.provider.pool.health_policy import apply_health_policy
|
||||
|
||||
await apply_health_policy(
|
||||
provider_id=self.provider_id,
|
||||
key_id=key_id,
|
||||
status_code=status_code,
|
||||
error_body=error_body,
|
||||
response_headers=response_headers,
|
||||
config=self.config,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Key schedulability check (used by candidate_builder)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def is_key_schedulable(self, key_id: str) -> tuple[bool, str | None]:
|
||||
"""Check if *key_id* is currently schedulable (not in cooldown, not
|
||||
cost-exhausted). Returns ``(True, None)`` or ``(False, reason)``.
|
||||
"""
|
||||
pid = self.provider_id
|
||||
|
||||
# Cooldown check
|
||||
cd = await redis_ops.get_cooldown(pid, key_id)
|
||||
if cd is not None:
|
||||
return False, f"pool cooldown: {cd}"
|
||||
|
||||
# Cost check
|
||||
if self.config.cost_limit_per_key_tokens is not None:
|
||||
total = await redis_ops.get_cost_window_total(
|
||||
pid, key_id, self.config.cost_window_seconds
|
||||
)
|
||||
if total >= self.config.cost_limit_per_key_tokens:
|
||||
return False, "pool cost limit reached"
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
# Backward-compatible alias
|
||||
ClaudeCodePoolManager = PoolManager
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
def _attach_pool_extra(candidate: Any, ct: PoolCandidateTrace) -> None:
|
||||
"""Attach pool trace extra_data onto a candidate object."""
|
||||
existing = getattr(candidate, "_pool_extra_data", None) or {}
|
||||
existing.update(ct.to_extra_data())
|
||||
setattr(candidate, "_pool_extra_data", existing)
|
||||
|
||||
|
||||
def _get_active_strategies(config: PoolConfig) -> list[Any]:
|
||||
"""Get active strategies for the given config (lazy import)."""
|
||||
if not config.strategies:
|
||||
return []
|
||||
try:
|
||||
# Import triggers built-in strategy registration via module-level side effects.
|
||||
from src.services.provider.pool import strategies as _builtin_strategies # noqa: F401
|
||||
from src.services.provider.pool.strategy import get_active_strategies
|
||||
|
||||
return get_active_strategies(config.strategies)
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def _shuffle_same_score(
|
||||
items: list[_T],
|
||||
lru_scores: dict[str, float],
|
||||
key_fn: Callable[[_T], str],
|
||||
) -> None:
|
||||
"""In-place random shuffle within groups that share the same LRU score."""
|
||||
if len(items) <= 1:
|
||||
return
|
||||
|
||||
i = 0
|
||||
while i < len(items):
|
||||
score_i = lru_scores.get(key_fn(items[i]), 0.0)
|
||||
j = i + 1
|
||||
while j < len(items) and lru_scores.get(key_fn(items[j]), 0.0) == score_i:
|
||||
j += 1
|
||||
if j - i > 1:
|
||||
group = items[i:j]
|
||||
random.shuffle(group)
|
||||
items[i:j] = group
|
||||
i = j
|
||||
|
||||
|
||||
def _shuffle_same_score_groups(
|
||||
candidates: list[ProviderCandidate],
|
||||
lru_scores: dict[str, float],
|
||||
) -> None:
|
||||
_shuffle_same_score(candidates, lru_scores, lambda c: str(c.key.id))
|
||||
42
_deprecated_py_src/services/provider/pool/oauth_cache.py
Normal file
42
_deprecated_py_src/services/provider/pool/oauth_cache.py
Normal file
@@ -0,0 +1,42 @@
|
||||
"""OAuth token Redis cache for the Account Pool.
|
||||
|
||||
Additions over the base ``auth.py`` refresh flow:
|
||||
|
||||
- **Redis token cache**: Avoids repeated DB decryption for hot keys.
|
||||
Cache key: ``provider_oauth_token_cache:{key_id}``
|
||||
- **Configurable proactive refresh skew**: Default 180 s (3 min) instead
|
||||
of the base 120 s, configurable via ``PoolConfig.proactive_refresh_seconds``.
|
||||
- **401 immediate invalidation**: Clears the Redis cache so the next request
|
||||
triggers a fresh refresh.
|
||||
|
||||
This module does NOT replace ``auth.py``; it adds a caching layer that
|
||||
``auth.py`` can consult before decrypting from DB.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.pool import redis_ops
|
||||
|
||||
|
||||
async def get_cached_token(key_id: str) -> str | None:
|
||||
"""Return cached access token from Redis, or None."""
|
||||
return await redis_ops.get_cached_oauth_token(key_id)
|
||||
|
||||
|
||||
async def cache_token(key_id: str, token: str, expires_in_seconds: int) -> None:
|
||||
"""Cache an access token in Redis.
|
||||
|
||||
*expires_in_seconds* is the remaining lifetime of the token. We shave
|
||||
off 60 s so the cache expires slightly before the token itself, giving
|
||||
the refresh flow time to act.
|
||||
"""
|
||||
ttl = max(1, expires_in_seconds - 60)
|
||||
await redis_ops.cache_oauth_token(key_id, token, ttl)
|
||||
logger.debug("Pool OAuth: cached token for key {} (TTL={}s)", key_id[:8], ttl)
|
||||
|
||||
|
||||
async def invalidate_token(key_id: str) -> None:
|
||||
"""Invalidate the cached token (e.g. after a 401)."""
|
||||
await redis_ops.invalidate_oauth_token_cache(key_id)
|
||||
logger.debug("Pool OAuth: invalidated token cache for key {}", key_id[:8])
|
||||
694
_deprecated_py_src/services/provider/pool/redis_ops.py
Normal file
694
_deprecated_py_src/services/provider/pool/redis_ops.py
Normal file
@@ -0,0 +1,694 @@
|
||||
"""Redis operations for the Account Pool (provider-agnostic).
|
||||
|
||||
All pool transient state is stored in Redis. This module centralises key
|
||||
naming, Lua scripts, and graceful fallbacks so that the rest of the pool
|
||||
layer is free of Redis specifics.
|
||||
|
||||
Key schema
|
||||
----------
|
||||
ap:{pid}:sticky:{session_uuid} STRING -> key_id (TTL: config)
|
||||
ap:{pid}:lru ZSET member=key_id, score=unix_ts
|
||||
ap:{pid}:cooldown:{key_id} STRING -> reason (TTL: error-specific)
|
||||
ap:{pid}:cost:{key_id} ZSET member=req_id, score=unix_ts
|
||||
ap:{pid}:latency:{key_id} ZSET member=req_id:ttfb_ms, score=unix_ts
|
||||
provider_oauth_token_cache:{key_id} STRING -> access_token (TTL: expires - 60)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from src.clients.redis_client import get_redis_client
|
||||
from src.core.logger import logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
PREFIX = "ap"
|
||||
|
||||
|
||||
def _sticky_key(provider_id: str, session_uuid: str) -> str:
|
||||
return f"{PREFIX}:{provider_id}:sticky:{session_uuid}"
|
||||
|
||||
|
||||
def _lru_key(provider_id: str) -> str:
|
||||
return f"{PREFIX}:{provider_id}:lru"
|
||||
|
||||
|
||||
def _cooldown_key(provider_id: str, key_id: str) -> str:
|
||||
return f"{PREFIX}:{provider_id}:cooldown:{key_id}"
|
||||
|
||||
|
||||
def _cooldown_index_key(provider_id: str) -> str:
|
||||
"""SET tracking which keys are in cooldown (for O(1) count queries)."""
|
||||
return f"{PREFIX}:{provider_id}:cooldown_idx"
|
||||
|
||||
|
||||
def _cost_key(provider_id: str, key_id: str) -> str:
|
||||
return f"{PREFIX}:{provider_id}:cost:{key_id}"
|
||||
|
||||
|
||||
def _latency_key(provider_id: str, key_id: str) -> str:
|
||||
return f"{PREFIX}:{provider_id}:latency:{key_id}"
|
||||
|
||||
|
||||
def _oauth_cache_key(key_id: str) -> str:
|
||||
return f"provider_oauth_token_cache:{key_id}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Lua scripts
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Sticky-select: GET binding, verify it's not in cooldown, refresh TTL.
|
||||
# KEYS[1] = sticky key, KEYS[2] = cooldown key prefix (ap:{pid}:cooldown:)
|
||||
# ARGV[1] = ttl
|
||||
# Returns: key_id or nil
|
||||
_STICKY_SELECT_LUA = """
|
||||
local binding = redis.call("GET", KEYS[1])
|
||||
if not binding then
|
||||
return nil
|
||||
end
|
||||
-- Check cooldown for the bound key
|
||||
local cooldown_key = KEYS[2] .. binding
|
||||
local in_cooldown = redis.call("EXISTS", cooldown_key)
|
||||
if in_cooldown == 1 then
|
||||
redis.call("DEL", KEYS[1])
|
||||
return nil
|
||||
end
|
||||
redis.call("EXPIRE", KEYS[1], tonumber(ARGV[1]))
|
||||
return binding
|
||||
"""
|
||||
|
||||
# Cost window sum (read-only, no cleanup on read path for performance).
|
||||
# KEYS[1] = cost zset key, ARGV[1] = window_start timestamp
|
||||
# Returns total token count within the window.
|
||||
_COST_WINDOW_SUM_LUA = """
|
||||
local key = KEYS[1]
|
||||
local window_start = tonumber(ARGV[1])
|
||||
local members = redis.call("ZRANGEBYSCORE", key, window_start, "+inf")
|
||||
local total = 0
|
||||
for _, m in ipairs(members) do
|
||||
local colon = string.find(m, ":", 1, true)
|
||||
if colon then
|
||||
local n = tonumber(string.sub(m, colon + 1))
|
||||
if n then total = total + n end
|
||||
end
|
||||
end
|
||||
return total
|
||||
"""
|
||||
|
||||
# Latency window average (read-only, no cleanup on read path for performance).
|
||||
# KEYS[1] = latency zset key, ARGV[1] = window_start timestamp
|
||||
# Returns nil when there are no samples, or avg(ms) as number.
|
||||
_LATENCY_WINDOW_AVG_LUA = """
|
||||
local key = KEYS[1]
|
||||
local window_start = tonumber(ARGV[1])
|
||||
local members = redis.call("ZRANGEBYSCORE", key, window_start, "+inf")
|
||||
local total = 0
|
||||
local count = 0
|
||||
for _, m in ipairs(members) do
|
||||
local colon = string.find(m, ":", 1, true)
|
||||
if colon then
|
||||
local n = tonumber(string.sub(m, colon + 1))
|
||||
if n then
|
||||
total = total + n
|
||||
count = count + 1
|
||||
end
|
||||
end
|
||||
end
|
||||
if count == 0 then
|
||||
return nil
|
||||
end
|
||||
return total / count
|
||||
"""
|
||||
|
||||
|
||||
async def _get_redis() -> "aioredis.Redis | None":
|
||||
return await get_redis_client(require_redis=False)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sticky session
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def get_sticky_binding(provider_id: str, session_uuid: str, ttl: int) -> str | None:
|
||||
"""Get and refresh sticky session binding. Returns key_id or None."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return None
|
||||
try:
|
||||
result = await redis.eval(
|
||||
_STICKY_SELECT_LUA,
|
||||
2,
|
||||
_sticky_key(provider_id, session_uuid),
|
||||
f"{PREFIX}:{provider_id}:cooldown:",
|
||||
str(ttl),
|
||||
)
|
||||
if result:
|
||||
return result.decode() if isinstance(result, bytes) else str(result)
|
||||
return None
|
||||
except Exception:
|
||||
logger.debug("Pool: sticky GET failed for session {}", session_uuid[:8])
|
||||
return None
|
||||
|
||||
|
||||
async def set_sticky_binding(provider_id: str, session_uuid: str, key_id: str, ttl: int) -> None:
|
||||
"""Create or update sticky session binding."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.setex(_sticky_key(provider_id, session_uuid), ttl, key_id)
|
||||
except Exception:
|
||||
logger.debug("Pool: sticky SET failed for session {}", session_uuid[:8])
|
||||
|
||||
|
||||
async def delete_sticky_binding(provider_id: str, session_uuid: str) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.delete(_sticky_key(provider_id, session_uuid))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# LRU
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def get_lru_scores(provider_id: str, key_ids: list[str]) -> dict[str, float]:
|
||||
"""Batch-fetch LRU timestamps. Missing keys get score 0 (highest priority)."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return {}
|
||||
try:
|
||||
lru_k = _lru_key(provider_id)
|
||||
scores = await redis.zmscore(lru_k, key_ids)
|
||||
result: dict[str, float] = {}
|
||||
for kid, score in zip(key_ids, scores):
|
||||
result[kid] = float(score) if score is not None else 0.0
|
||||
return result
|
||||
except Exception:
|
||||
logger.debug("Pool: LRU ZMSCORE failed for provider {}", provider_id[:8])
|
||||
return {}
|
||||
|
||||
|
||||
async def touch_lru(provider_id: str, key_id: str) -> None:
|
||||
"""Update last-used timestamp."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.zadd(_lru_key(provider_id), {key_id: time.time()})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cooldown
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def set_cooldown(provider_id: str, key_id: str, reason: str, ttl: int) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
pipe = redis.pipeline()
|
||||
pipe.setex(_cooldown_key(provider_id, key_id), ttl, reason)
|
||||
# Track in index set for O(1) count queries.
|
||||
idx_key = _cooldown_index_key(provider_id)
|
||||
pipe.sadd(idx_key, key_id)
|
||||
# Keep index alive at least as long as the longest cooldown entry.
|
||||
# Each set_cooldown call refreshes the TTL so the SET won't expire
|
||||
# while there are still active cooldowns.
|
||||
pipe.expire(idx_key, ttl + 60)
|
||||
await pipe.execute()
|
||||
logger.info(
|
||||
"Pool: key {} cooldown set: {} ({}s)",
|
||||
key_id[:8],
|
||||
reason,
|
||||
ttl,
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("Pool: cooldown SET failed for key {}", key_id[:8])
|
||||
|
||||
|
||||
async def get_cooldown(provider_id: str, key_id: str) -> str | None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return None
|
||||
try:
|
||||
val = await redis.get(_cooldown_key(provider_id, key_id))
|
||||
if val:
|
||||
return val.decode() if isinstance(val, bytes) else str(val)
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def clear_cooldown(provider_id: str, key_id: str) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
pipe = redis.pipeline()
|
||||
pipe.delete(_cooldown_key(provider_id, key_id))
|
||||
pipe.srem(_cooldown_index_key(provider_id), key_id)
|
||||
await pipe.execute()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
async def batch_get_cooldowns(
|
||||
provider_id: str,
|
||||
key_ids: list[str],
|
||||
*,
|
||||
include_ttl: bool = False,
|
||||
) -> dict[str, str | None] | dict[str, tuple[str | None, int | None]]:
|
||||
"""Batch check cooldown status for multiple keys.
|
||||
|
||||
When *include_ttl* is ``True``, each value is a ``(reason, ttl_seconds)``
|
||||
tuple instead of a plain reason string. The TTL commands are batched in
|
||||
the same pipeline so there is no extra round-trip.
|
||||
"""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
if include_ttl:
|
||||
return {k: (None, None) for k in key_ids}
|
||||
return {k: None for k in key_ids}
|
||||
try:
|
||||
pipe = redis.pipeline()
|
||||
for kid in key_ids:
|
||||
ck = _cooldown_key(provider_id, kid)
|
||||
pipe.get(ck)
|
||||
if include_ttl:
|
||||
pipe.ttl(ck)
|
||||
results = await pipe.execute()
|
||||
|
||||
if include_ttl:
|
||||
out_ttl: dict[str, tuple[str | None, int | None]] = {}
|
||||
# results interleave GET/TTL: [val0, ttl0, val1, ttl1, ...]
|
||||
for i, kid in enumerate(key_ids):
|
||||
val = results[i * 2]
|
||||
ttl_val = results[i * 2 + 1]
|
||||
reason: str | None = None
|
||||
if val:
|
||||
reason = val.decode() if isinstance(val, bytes) else str(val)
|
||||
ttl_sec: int | None = None
|
||||
if isinstance(ttl_val, int) and ttl_val > 0:
|
||||
ttl_sec = ttl_val
|
||||
out_ttl[kid] = (reason, ttl_sec)
|
||||
return out_ttl
|
||||
|
||||
out: dict[str, str | None] = {}
|
||||
for kid, val in zip(key_ids, results):
|
||||
if val:
|
||||
out[kid] = val.decode() if isinstance(val, bytes) else str(val)
|
||||
else:
|
||||
out[kid] = None
|
||||
return out
|
||||
except Exception:
|
||||
if include_ttl:
|
||||
return {k: (None, None) for k in key_ids}
|
||||
return {k: None for k in key_ids}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cost tracking
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def add_cost_entry(provider_id: str, key_id: str, tokens: int, window_seconds: int) -> None:
|
||||
"""Record a cost entry (tokens used) with automatic window expiry."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
now = time.time()
|
||||
cost_k = _cost_key(provider_id, key_id)
|
||||
member = f"{uuid.uuid4().hex}:{tokens}"
|
||||
window_start = now - max(int(window_seconds), 1)
|
||||
pipe = redis.pipeline()
|
||||
pipe.zadd(cost_k, {member: now})
|
||||
# Prune expired entries on the write path (moved from read Lua script).
|
||||
pipe.zremrangebyscore(cost_k, "-inf", window_start)
|
||||
# Set a TTL slightly larger than the window to auto-clean abandoned keys.
|
||||
pipe.expire(cost_k, window_seconds + 600)
|
||||
await pipe.execute()
|
||||
except Exception:
|
||||
logger.debug("Pool: cost ADD failed for key {}", key_id[:8])
|
||||
|
||||
|
||||
async def get_cost_window_total(provider_id: str, key_id: str, window_seconds: int) -> int:
|
||||
"""Sum tokens used within the rolling window (single key)."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return 0
|
||||
try:
|
||||
now = time.time()
|
||||
window_start = now - window_seconds
|
||||
cost_k = _cost_key(provider_id, key_id)
|
||||
result = await redis.eval(_COST_WINDOW_SUM_LUA, 1, cost_k, str(window_start))
|
||||
return int(result) if result else 0
|
||||
except Exception:
|
||||
logger.debug("Pool: cost SUM failed for key {}", key_id[:8])
|
||||
return 0
|
||||
|
||||
|
||||
async def batch_get_cost_totals(
|
||||
provider_id: str, key_ids: list[str], window_seconds: int
|
||||
) -> dict[str, int]:
|
||||
"""Batch-fetch cost totals for multiple keys using pipeline + Lua."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return {k: 0 for k in key_ids}
|
||||
try:
|
||||
now = time.time()
|
||||
window_start = now - window_seconds
|
||||
pipe = redis.pipeline()
|
||||
for kid in key_ids:
|
||||
cost_k = _cost_key(provider_id, kid)
|
||||
pipe.eval(_COST_WINDOW_SUM_LUA, 1, cost_k, str(window_start))
|
||||
results = await pipe.execute()
|
||||
out: dict[str, int] = {}
|
||||
for kid, val in zip(key_ids, results):
|
||||
out[kid] = int(val) if val else 0
|
||||
return out
|
||||
except Exception:
|
||||
logger.debug("Pool: batch cost SUM failed for provider {}", provider_id[:8])
|
||||
return {k: 0 for k in key_ids}
|
||||
|
||||
|
||||
async def record_latency(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
ttfb_ms: int,
|
||||
window_seconds: int,
|
||||
sample_limit: int,
|
||||
) -> None:
|
||||
"""Record one TTFB sample with rolling-window cleanup."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
now = time.time()
|
||||
latency_k = _latency_key(provider_id, key_id)
|
||||
sample = max(int(ttfb_ms), 0)
|
||||
member = f"{uuid.uuid4().hex}:{sample}"
|
||||
pipe = redis.pipeline()
|
||||
pipe.zadd(latency_k, {member: now})
|
||||
window_start = now - max(int(window_seconds), 1)
|
||||
pipe.zremrangebyscore(latency_k, "-inf", window_start)
|
||||
capped_limit = max(int(sample_limit), 1)
|
||||
pipe.zremrangebyrank(latency_k, 0, -(capped_limit + 1))
|
||||
pipe.expire(latency_k, max(int(window_seconds), 1) + 600)
|
||||
await pipe.execute()
|
||||
except Exception:
|
||||
logger.debug("Pool: latency ADD failed for key {}", key_id[:8])
|
||||
|
||||
|
||||
async def batch_get_latency_avgs(
|
||||
provider_id: str,
|
||||
key_ids: list[str],
|
||||
window_seconds: int,
|
||||
) -> dict[str, float]:
|
||||
"""Batch-fetch latency averages (ms) for keys in a rolling window."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return {}
|
||||
try:
|
||||
now = time.time()
|
||||
window_start = now - max(int(window_seconds), 1)
|
||||
pipe = redis.pipeline()
|
||||
for kid in key_ids:
|
||||
pipe.eval(_LATENCY_WINDOW_AVG_LUA, 1, _latency_key(provider_id, kid), str(window_start))
|
||||
results = await pipe.execute()
|
||||
out: dict[str, float] = {}
|
||||
for kid, val in zip(key_ids, results):
|
||||
if val is None:
|
||||
continue
|
||||
try:
|
||||
out[kid] = float(val)
|
||||
except Exception:
|
||||
continue
|
||||
return out
|
||||
except Exception:
|
||||
logger.debug("Pool: batch latency AVG failed for provider {}", provider_id[:8])
|
||||
return {}
|
||||
|
||||
|
||||
async def clear_cost(provider_id: str, key_id: str) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.delete(_cost_key(provider_id, key_id))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OAuth token cache
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def cache_oauth_token(key_id: str, token: str, ttl: int) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
if ttl > 0:
|
||||
await redis.setex(_oauth_cache_key(key_id), ttl, token)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
async def get_cached_oauth_token(key_id: str) -> str | None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return None
|
||||
try:
|
||||
val = await redis.get(_oauth_cache_key(key_id))
|
||||
if val:
|
||||
return val.decode() if isinstance(val, bytes) else str(val)
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def invalidate_oauth_token_cache(key_id: str) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.delete(_oauth_cache_key(key_id))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pool status query (admin)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def get_sticky_session_count(provider_id: str) -> int:
|
||||
"""Approximate count of active sticky sessions (via SCAN, for admin display only)."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return 0
|
||||
try:
|
||||
pattern = f"{PREFIX}:{provider_id}:sticky:*"
|
||||
count = 0
|
||||
async for _ in redis.scan_iter(match=pattern, count=100):
|
||||
count += 1
|
||||
return count
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
|
||||
async def get_key_sticky_count(provider_id: str, key_id: str) -> int:
|
||||
"""Count sticky sessions bound to a specific key (admin only).
|
||||
|
||||
Uses batched SCAN + pipeline MGET to reduce Redis round-trips.
|
||||
"""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return 0
|
||||
try:
|
||||
pattern = f"{PREFIX}:{provider_id}:sticky:*"
|
||||
count = 0
|
||||
batch: list[bytes | str] = []
|
||||
async for k in redis.scan_iter(match=pattern, count=200):
|
||||
batch.append(k)
|
||||
if len(batch) >= 200:
|
||||
vals = await redis.mget(batch)
|
||||
for val in vals:
|
||||
if val:
|
||||
bound_id = val.decode() if isinstance(val, bytes) else str(val)
|
||||
if bound_id == key_id:
|
||||
count += 1
|
||||
batch.clear()
|
||||
if batch:
|
||||
vals = await redis.mget(batch)
|
||||
for val in vals:
|
||||
if val:
|
||||
bound_id = val.decode() if isinstance(val, bytes) else str(val)
|
||||
if bound_id == key_id:
|
||||
count += 1
|
||||
return count
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
|
||||
async def batch_get_key_sticky_counts(
|
||||
provider_id: str,
|
||||
key_ids: list[str],
|
||||
) -> dict[str, int]:
|
||||
"""Count sticky sessions for multiple keys in a single scan (admin only)."""
|
||||
if not key_ids:
|
||||
return {}
|
||||
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return {kid: 0 for kid in key_ids}
|
||||
|
||||
target_ids = set(key_ids)
|
||||
counts: dict[str, int] = {kid: 0 for kid in key_ids}
|
||||
|
||||
try:
|
||||
pattern = f"{PREFIX}:{provider_id}:sticky:*"
|
||||
batch: list[bytes | str] = []
|
||||
async for key in redis.scan_iter(match=pattern, count=200):
|
||||
batch.append(key)
|
||||
if len(batch) >= 200:
|
||||
vals = await redis.mget(batch)
|
||||
for val in vals:
|
||||
if not val:
|
||||
continue
|
||||
bound_id = val.decode() if isinstance(val, bytes) else str(val)
|
||||
if bound_id in target_ids:
|
||||
counts[bound_id] = counts.get(bound_id, 0) + 1
|
||||
batch.clear()
|
||||
|
||||
if batch:
|
||||
vals = await redis.mget(batch)
|
||||
for val in vals:
|
||||
if not val:
|
||||
continue
|
||||
bound_id = val.decode() if isinstance(val, bytes) else str(val)
|
||||
if bound_id in target_ids:
|
||||
counts[bound_id] = counts.get(bound_id, 0) + 1
|
||||
|
||||
return counts
|
||||
except Exception:
|
||||
return {kid: 0 for kid in key_ids}
|
||||
|
||||
|
||||
async def get_cooldown_ttl(provider_id: str, key_id: str) -> int | None:
|
||||
"""Get remaining cooldown TTL in seconds. None = no cooldown."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return None
|
||||
try:
|
||||
ttl = await redis.ttl(_cooldown_key(provider_id, key_id))
|
||||
return ttl if ttl > 0 else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def batch_get_cooldown_ttls(provider_id: str, key_ids: list[str]) -> dict[str, int | None]:
|
||||
"""Batch-fetch cooldown TTLs for multiple keys using pipeline."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return {k: None for k in key_ids}
|
||||
try:
|
||||
pipe = redis.pipeline()
|
||||
for kid in key_ids:
|
||||
pipe.ttl(_cooldown_key(provider_id, kid))
|
||||
results = await pipe.execute()
|
||||
out: dict[str, int | None] = {}
|
||||
for kid, ttl in zip(key_ids, results):
|
||||
out[kid] = int(ttl) if isinstance(ttl, int) and ttl > 0 else None
|
||||
return out
|
||||
except Exception:
|
||||
return {k: None for k in key_ids}
|
||||
|
||||
|
||||
async def batch_count_provider_cooldowns(provider_ids: list[str]) -> dict[str, int]:
|
||||
"""Count cooldown entries per provider using the cooldown index set.
|
||||
|
||||
Uses ``SCARD`` on the ``ap:{pid}:cooldown_idx`` set for O(1) count
|
||||
instead of scanning the key-space. The index set is maintained by
|
||||
:func:`set_cooldown` / :func:`clear_cooldown`.
|
||||
|
||||
Note: the index set may contain stale entries (expired cooldowns whose
|
||||
TTL elapsed before an explicit clear). This over-count is acceptable
|
||||
for admin display purposes -- precision is not critical here.
|
||||
"""
|
||||
if not provider_ids:
|
||||
return {}
|
||||
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return {pid: 0 for pid in provider_ids}
|
||||
|
||||
try:
|
||||
pipe = redis.pipeline()
|
||||
for pid in provider_ids:
|
||||
pipe.scard(_cooldown_index_key(pid))
|
||||
results = await pipe.execute()
|
||||
counts: dict[str, int] = {}
|
||||
for pid, val in zip(provider_ids, results):
|
||||
counts[pid] = max(int(val or 0), 0)
|
||||
return counts
|
||||
except Exception:
|
||||
return {pid: 0 for pid in provider_ids}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Stream timeout counter
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_STREAM_TIMEOUT_KEY_FMT = f"{PREFIX}:{{}}:stream_timeout:{{}}"
|
||||
|
||||
|
||||
def _stream_timeout_key(provider_id: str, key_id: str) -> str:
|
||||
return _STREAM_TIMEOUT_KEY_FMT.format(provider_id, key_id)
|
||||
|
||||
|
||||
async def incr_stream_timeout_count(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
window_seconds: int,
|
||||
) -> int:
|
||||
"""Increment stream timeout counter and return count within the window.
|
||||
|
||||
Uses a ZSET with timestamps as scores. Old entries beyond the window
|
||||
are pruned on each call. Returns the count of timeouts in the window.
|
||||
"""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return 0
|
||||
try:
|
||||
now = time.time()
|
||||
window_start = now - window_seconds
|
||||
key = _stream_timeout_key(provider_id, key_id)
|
||||
member = f"{uuid.uuid4().hex}"
|
||||
pipe = redis.pipeline()
|
||||
pipe.zremrangebyscore(key, "-inf", window_start)
|
||||
pipe.zadd(key, {member: now})
|
||||
pipe.zcard(key)
|
||||
pipe.expire(key, window_seconds + 60)
|
||||
results = await pipe.execute()
|
||||
count = int(results[2]) if results[2] else 0
|
||||
return count
|
||||
except Exception:
|
||||
logger.debug("Pool: stream timeout INCR failed for key {}", key_id[:8])
|
||||
return 0
|
||||
@@ -0,0 +1,466 @@
|
||||
"""Pool scheduling dimension registry and evaluation helpers.
|
||||
|
||||
This module keeps pool scheduling scoring isolated from API layer code.
|
||||
Callers build a :class:`PoolSchedulingSnapshot` and evaluate it against
|
||||
registered dimensions to obtain a normalized summary.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Protocol
|
||||
|
||||
PoolDimensionStatus = str # ok / degraded / blocked
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PoolSchedulingSnapshot:
|
||||
"""Point-in-time scheduling inputs for one key."""
|
||||
|
||||
is_active: bool
|
||||
cooldown_reason: str | None
|
||||
cooldown_ttl_seconds: int | None
|
||||
circuit_breaker_open: bool
|
||||
cost_window_usage: int
|
||||
cost_limit: int | None
|
||||
cost_soft_threshold_percent: int = 80
|
||||
health_score: float = 1.0
|
||||
latency_avg_ms: float | None = None
|
||||
account_blocked: bool = False
|
||||
account_block_label: str | None = None
|
||||
account_block_reason: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PoolSchedulingDimensionResult:
|
||||
"""Evaluation output for one scheduling dimension."""
|
||||
|
||||
code: str
|
||||
label: str
|
||||
status: PoolDimensionStatus = "ok"
|
||||
blocking: bool = False
|
||||
source: str = "pool"
|
||||
weight: int = 1
|
||||
score: float = 1.0
|
||||
detail: str | None = None
|
||||
ttl_seconds: int | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PoolSchedulingSummary:
|
||||
"""Merged scheduling state across all dimensions."""
|
||||
|
||||
status: str # available / degraded / blocked
|
||||
reason: str
|
||||
label: str
|
||||
candidate_eligible: bool
|
||||
blocked_count: int
|
||||
degraded_count: int
|
||||
|
||||
|
||||
class PoolSchedulingDimension(Protocol):
|
||||
"""Dimension evaluator protocol."""
|
||||
|
||||
code: str
|
||||
label: str
|
||||
source: str
|
||||
weight: int
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
"""Evaluate one dimension from snapshot."""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _AccountStateDimension:
|
||||
code: str = "account_state"
|
||||
label: str = "账号状态"
|
||||
source: str = "policy"
|
||||
weight: int = 10
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
if not snapshot.account_blocked:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
)
|
||||
|
||||
blocked_label = snapshot.account_block_label or "账号异常"
|
||||
if blocked_label == "账号封禁":
|
||||
blocked_code = "account_banned"
|
||||
elif blocked_label == "工作区停用":
|
||||
blocked_code = "workspace_deactivated"
|
||||
elif blocked_label == "账号停用":
|
||||
blocked_code = "account_disabled"
|
||||
elif blocked_label == "需要验证":
|
||||
blocked_code = "account_verification"
|
||||
elif blocked_label == "访问受限":
|
||||
blocked_code = "account_forbidden"
|
||||
else:
|
||||
blocked_code = "account_blocked"
|
||||
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=blocked_code,
|
||||
label=blocked_label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="blocked",
|
||||
blocking=True,
|
||||
score=0.0,
|
||||
detail=snapshot.account_block_reason or blocked_label,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ManualEnableDimension:
|
||||
code: str = "manual_disabled"
|
||||
label: str = "禁用"
|
||||
source: str = "manual"
|
||||
weight: int = 8
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
if snapshot.is_active:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
)
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="blocked",
|
||||
blocking=True,
|
||||
score=0.0,
|
||||
detail="账号被手动禁用",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CooldownDimension:
|
||||
code: str = "cooldown"
|
||||
label: str = "冷却中"
|
||||
source: str = "pool"
|
||||
weight: int = 7
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
if not snapshot.cooldown_reason:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
)
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="blocked",
|
||||
blocking=True,
|
||||
score=0.0,
|
||||
detail=snapshot.cooldown_reason,
|
||||
ttl_seconds=snapshot.cooldown_ttl_seconds,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CircuitBreakerDimension:
|
||||
code: str = "circuit_open"
|
||||
label: str = "熔断中"
|
||||
source: str = "health"
|
||||
weight: int = 6
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
if not snapshot.circuit_breaker_open:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
)
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="blocked",
|
||||
blocking=True,
|
||||
score=0.0,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CostDimension:
|
||||
code: str = "cost"
|
||||
label: str = "成本"
|
||||
source: str = "pool"
|
||||
weight: int = 5
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
limit = snapshot.cost_limit
|
||||
usage = max(snapshot.cost_window_usage, 0)
|
||||
if limit is None or limit <= 0:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
detail=f"{usage}/-",
|
||||
)
|
||||
|
||||
ratio = usage / limit
|
||||
detail = f"{usage}/{limit}"
|
||||
if ratio >= 1.0:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code="cost_exhausted",
|
||||
label="成本超限",
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="blocked",
|
||||
blocking=True,
|
||||
score=0.0,
|
||||
detail=detail,
|
||||
)
|
||||
|
||||
soft_threshold = max(1, min(snapshot.cost_soft_threshold_percent, 100))
|
||||
if ratio * 100 >= soft_threshold:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code="cost_soft",
|
||||
label="成本接近上限",
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="degraded",
|
||||
score=0.45,
|
||||
detail=detail,
|
||||
)
|
||||
|
||||
if ratio >= 0.6:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="degraded",
|
||||
score=0.72,
|
||||
detail=detail,
|
||||
)
|
||||
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
detail=detail,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _HealthDimension:
|
||||
code: str = "health"
|
||||
label: str = "健康度"
|
||||
source: str = "health"
|
||||
weight: int = 4
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
score = max(0.0, min(snapshot.health_score, 1.0))
|
||||
detail = f"{score:.2f}"
|
||||
if score < 0.5:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code="health_low",
|
||||
label="健康度过低",
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="degraded",
|
||||
score=0.3,
|
||||
detail=detail,
|
||||
)
|
||||
if score < 0.8:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code="health_degraded",
|
||||
label="健康度下降",
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="degraded",
|
||||
score=0.65,
|
||||
detail=detail,
|
||||
)
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
detail=detail,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _LatencyDimension:
|
||||
code: str = "latency"
|
||||
label: str = "延迟"
|
||||
source: str = "runtime"
|
||||
weight: int = 3
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
latency = snapshot.latency_avg_ms
|
||||
if latency is None:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
detail="-",
|
||||
)
|
||||
|
||||
value = max(float(latency), 0.0)
|
||||
detail = f"{value:.0f}ms"
|
||||
if value >= 3000:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code="latency_high",
|
||||
label="延迟偏高",
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="degraded",
|
||||
score=0.5,
|
||||
detail=detail,
|
||||
)
|
||||
if value >= 1200:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code="latency_slow",
|
||||
label="延迟较慢",
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="degraded",
|
||||
score=0.72,
|
||||
detail=detail,
|
||||
)
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
detail=detail,
|
||||
)
|
||||
|
||||
|
||||
_POOL_DIMENSION_REGISTRY: dict[str, PoolSchedulingDimension] = {}
|
||||
_POOL_DIMENSION_ORDER: list[str] = []
|
||||
|
||||
|
||||
def register_pool_scheduling_dimension(name: str, dimension: PoolSchedulingDimension) -> None:
|
||||
"""Register a dimension evaluator by name."""
|
||||
normalized = name.strip()
|
||||
if not normalized:
|
||||
return
|
||||
if normalized not in _POOL_DIMENSION_ORDER:
|
||||
_POOL_DIMENSION_ORDER.append(normalized)
|
||||
_POOL_DIMENSION_REGISTRY[normalized] = dimension
|
||||
|
||||
|
||||
def get_pool_scheduling_dimension(name: str) -> PoolSchedulingDimension | None:
|
||||
"""Fetch a registered dimension evaluator."""
|
||||
return _POOL_DIMENSION_REGISTRY.get(name.strip())
|
||||
|
||||
|
||||
def list_pool_scheduling_dimensions() -> tuple[str, ...]:
|
||||
"""List registered dimension names in evaluation order."""
|
||||
return tuple(_POOL_DIMENSION_ORDER)
|
||||
|
||||
|
||||
def evaluate_pool_scheduling_dimensions(
|
||||
snapshot: PoolSchedulingSnapshot,
|
||||
*,
|
||||
dimension_names: tuple[str, ...] | None = None,
|
||||
) -> list[PoolSchedulingDimensionResult]:
|
||||
"""Evaluate snapshot across all registered dimensions."""
|
||||
names = dimension_names or list_pool_scheduling_dimensions()
|
||||
results: list[PoolSchedulingDimensionResult] = []
|
||||
for name in names:
|
||||
dimension = get_pool_scheduling_dimension(name)
|
||||
if dimension is None:
|
||||
continue
|
||||
results.append(dimension.evaluate(snapshot))
|
||||
return results
|
||||
|
||||
|
||||
def summarize_pool_scheduling_dimensions(
|
||||
dimensions: list[PoolSchedulingDimensionResult],
|
||||
) -> PoolSchedulingSummary:
|
||||
"""Summarize dimension outputs into a unified scheduling state."""
|
||||
if not dimensions:
|
||||
return PoolSchedulingSummary(
|
||||
status="available",
|
||||
reason="available",
|
||||
label="可用",
|
||||
candidate_eligible=True,
|
||||
blocked_count=0,
|
||||
degraded_count=0,
|
||||
)
|
||||
|
||||
blocked = [item for item in dimensions if item.status == "blocked" or item.blocking]
|
||||
degraded = [item for item in dimensions if item.status == "degraded"]
|
||||
|
||||
if blocked:
|
||||
primary = blocked[0]
|
||||
return PoolSchedulingSummary(
|
||||
status="blocked",
|
||||
reason=primary.code,
|
||||
label=primary.label,
|
||||
candidate_eligible=False,
|
||||
blocked_count=len(blocked),
|
||||
degraded_count=len(degraded),
|
||||
)
|
||||
|
||||
if degraded:
|
||||
primary = degraded[0]
|
||||
return PoolSchedulingSummary(
|
||||
status="degraded",
|
||||
reason=primary.code,
|
||||
label=primary.label,
|
||||
candidate_eligible=True,
|
||||
blocked_count=0,
|
||||
degraded_count=len(degraded),
|
||||
)
|
||||
|
||||
return PoolSchedulingSummary(
|
||||
status="available",
|
||||
reason="available",
|
||||
label="可用",
|
||||
candidate_eligible=True,
|
||||
blocked_count=0,
|
||||
degraded_count=0,
|
||||
)
|
||||
|
||||
|
||||
def _register_default_dimensions() -> None:
|
||||
register_pool_scheduling_dimension("account_state", _AccountStateDimension())
|
||||
register_pool_scheduling_dimension("manual", _ManualEnableDimension())
|
||||
register_pool_scheduling_dimension("cooldown", _CooldownDimension())
|
||||
register_pool_scheduling_dimension("circuit", _CircuitBreakerDimension())
|
||||
register_pool_scheduling_dimension("cost", _CostDimension())
|
||||
register_pool_scheduling_dimension("latency", _LatencyDimension())
|
||||
register_pool_scheduling_dimension("health", _HealthDimension())
|
||||
|
||||
|
||||
_register_default_dimensions()
|
||||
@@ -0,0 +1,8 @@
|
||||
"""Built-in pool strategies."""
|
||||
|
||||
# Import side effects: register built-in strategies.
|
||||
import src.services.provider.pool.dimensions # noqa: F401
|
||||
|
||||
from . import multi_score # noqa: F401
|
||||
|
||||
__all__ = ["multi_score"]
|
||||
@@ -0,0 +1,264 @@
|
||||
"""Multi-dimension pool scoring strategy."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import src.services.provider.pool.dimensions # noqa: F401
|
||||
from src.services.provider.pool.dimensions import get_preset_dimension, get_preset_names
|
||||
from src.services.provider.pool.dimensions._helpers import rank_ascending, safe_float
|
||||
from src.services.provider.pool.strategy import register_pool_strategy
|
||||
|
||||
|
||||
def _normalize_mutex_group(value: Any) -> str | None:
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
normalized = value.strip().lower()
|
||||
return normalized or None
|
||||
|
||||
|
||||
def _get_preset_mutex_group(preset_name: str) -> str | None:
|
||||
# LRU is a built-in preset (not in registry) but shares the distribution mutex group.
|
||||
if preset_name == "lru":
|
||||
return "distribution_mode"
|
||||
dim = get_preset_dimension(preset_name)
|
||||
if dim is None:
|
||||
return None
|
||||
return _normalize_mutex_group(getattr(dim, "mutex_group", None))
|
||||
|
||||
|
||||
def _normalize_presets_from_config(
|
||||
config: Any,
|
||||
*,
|
||||
provider_type: str | None = None,
|
||||
) -> tuple[tuple[str, str | None], ...]:
|
||||
"""Extract enabled (preset_name, mode) tuples from config.scheduling_presets.
|
||||
|
||||
Supports both new SchedulingPreset objects and legacy string lists.
|
||||
Excludes ``lru`` from output (LRU is a final tie-breaker only).
|
||||
For mutex groups, enabled members inherit the group's first appearance index
|
||||
so the selected member keeps the group's visible priority slot.
|
||||
"""
|
||||
|
||||
raw = getattr(config, "scheduling_presets", ())
|
||||
if not isinstance(raw, (list, tuple)):
|
||||
return ()
|
||||
|
||||
normalized_provider_type = str(provider_type or "").strip().lower()
|
||||
allowed = get_preset_names() | {"lru"}
|
||||
entries: list[tuple[int, str, bool, str | None]] = []
|
||||
seen: set[str] = set()
|
||||
for idx, item in enumerate(raw):
|
||||
preset_name: str | None = None
|
||||
enabled = True
|
||||
mode: str | None = None
|
||||
|
||||
if hasattr(item, "preset"):
|
||||
preset_name = str(getattr(item, "preset", "")).strip().lower()
|
||||
enabled = bool(getattr(item, "enabled", True))
|
||||
raw_mode = getattr(item, "mode", None)
|
||||
if isinstance(raw_mode, str):
|
||||
mode = raw_mode.strip().lower() or None
|
||||
elif isinstance(item, str):
|
||||
preset_name = item.strip().lower()
|
||||
else:
|
||||
continue
|
||||
|
||||
if not preset_name or preset_name not in allowed or preset_name in seen:
|
||||
continue
|
||||
seen.add(preset_name)
|
||||
entries.append((idx, preset_name, enabled, mode))
|
||||
|
||||
# Codex 默认启用额度刷新优先维度(除非显式配置了 recent_refresh)。
|
||||
if (
|
||||
normalized_provider_type == "codex"
|
||||
and entries
|
||||
and "recent_refresh" not in {name for _idx, name, _enabled, _mode in entries}
|
||||
and "recent_refresh" in allowed
|
||||
):
|
||||
entries.append((len(entries), "recent_refresh", True, None))
|
||||
|
||||
if not entries:
|
||||
return ()
|
||||
|
||||
group_anchor_index: dict[str, int] = {}
|
||||
for idx, preset_name, _enabled, _mode in entries:
|
||||
mutex_group = _get_preset_mutex_group(preset_name)
|
||||
if mutex_group and mutex_group not in group_anchor_index:
|
||||
group_anchor_index[mutex_group] = idx
|
||||
|
||||
ordered_enabled: list[tuple[int, int, str, str | None]] = []
|
||||
group_enabled: dict[str, tuple[int, int, str, str | None]] = {}
|
||||
for idx, preset_name, enabled, mode in entries:
|
||||
if not enabled or preset_name == "lru":
|
||||
continue
|
||||
mutex_group = _get_preset_mutex_group(preset_name)
|
||||
if not mutex_group:
|
||||
ordered_enabled.append((idx, idx, preset_name, mode))
|
||||
continue
|
||||
|
||||
anchor = group_anchor_index.get(mutex_group, idx)
|
||||
existing = group_enabled.get(mutex_group)
|
||||
if existing is None or idx < existing[1]:
|
||||
group_enabled[mutex_group] = (anchor, idx, preset_name, mode)
|
||||
|
||||
ordered_enabled.extend(group_enabled.values())
|
||||
ordered_enabled.sort(key=lambda item: (item[0], item[1]))
|
||||
return tuple((preset_name, mode) for _anchor, _idx, preset_name, mode in ordered_enabled)
|
||||
|
||||
|
||||
class MultiScoreStrategy:
|
||||
name = "multi_score"
|
||||
|
||||
def compute_score(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
config: Any,
|
||||
context: dict[str, Any],
|
||||
) -> float | None:
|
||||
mode = str(getattr(config, "scheduling_mode", "lru") or "lru").strip().lower()
|
||||
if mode != "multi_score":
|
||||
return None
|
||||
|
||||
all_key_ids = [str(k) for k in (context.get("all_key_ids") or []) if str(k)]
|
||||
if not all_key_ids:
|
||||
return None
|
||||
|
||||
lru_scores = context.get("lru_scores", {})
|
||||
if not isinstance(lru_scores, dict):
|
||||
lru_scores = {}
|
||||
latency_avgs = context.get("latency_avgs", {})
|
||||
if not isinstance(latency_avgs, dict):
|
||||
latency_avgs = {}
|
||||
health_scores = context.get("health_scores", {})
|
||||
if not isinstance(health_scores, dict):
|
||||
health_scores = {}
|
||||
cost_totals = context.get("cost_totals", {})
|
||||
if not isinstance(cost_totals, dict):
|
||||
cost_totals = {}
|
||||
keys_by_id = context.get("keys_by_id", {})
|
||||
if not isinstance(keys_by_id, dict):
|
||||
keys_by_id = {}
|
||||
|
||||
presets = _normalize_presets_from_config(
|
||||
config,
|
||||
provider_type=context.get("provider_type"),
|
||||
)
|
||||
lru_enabled = bool(getattr(config, "lru_enabled", True))
|
||||
if presets:
|
||||
return self._compute_preset_score(
|
||||
key_id=key_id,
|
||||
all_key_ids=all_key_ids,
|
||||
presets=presets,
|
||||
lru_enabled=lru_enabled,
|
||||
lru_scores=lru_scores,
|
||||
keys_by_id=keys_by_id,
|
||||
context=context,
|
||||
)
|
||||
|
||||
weights = getattr(config, "scoring_weights", None)
|
||||
w_lru = safe_float(getattr(weights, "lru", 0.3)) or 0.0
|
||||
w_latency = safe_float(getattr(weights, "latency", 0.25)) or 0.0
|
||||
w_health = safe_float(getattr(weights, "health", 0.2)) or 0.0
|
||||
w_cost = safe_float(getattr(weights, "cost_remaining", 0.25)) or 0.0
|
||||
|
||||
lru_rank = rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
latency_rank = rank_ascending(key_id, latency_avgs, all_key_ids)
|
||||
|
||||
health_raw = safe_float(health_scores.get(key_id))
|
||||
if health_raw is None:
|
||||
health_raw = 1.0
|
||||
health_norm = 1.0 - max(0.0, min(health_raw, 1.0))
|
||||
|
||||
cost_limit = getattr(config, "cost_limit_per_key_tokens", None)
|
||||
used = safe_float(cost_totals.get(key_id)) or 0.0
|
||||
if cost_limit is None or int(cost_limit) <= 0:
|
||||
cost_norm = 0.0
|
||||
else:
|
||||
cost_norm = max(0.0, min(used / float(cost_limit), 1.0))
|
||||
|
||||
return (
|
||||
w_lru * lru_rank
|
||||
+ w_latency * latency_rank
|
||||
+ w_health * health_norm
|
||||
+ w_cost * cost_norm
|
||||
)
|
||||
|
||||
def _compute_preset_score(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
presets: tuple[tuple[str, str | None], ...],
|
||||
lru_enabled: bool,
|
||||
lru_scores: dict[str, Any],
|
||||
keys_by_id: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
) -> float:
|
||||
cache_signature = (tuple(all_key_ids), presets, bool(lru_enabled))
|
||||
cache = context.get("_preset_hard_order_cache")
|
||||
if (
|
||||
isinstance(cache, dict)
|
||||
and cache.get("signature") == cache_signature
|
||||
and isinstance(cache.get("ranks"), dict)
|
||||
):
|
||||
cached_rank = safe_float(cache["ranks"].get(key_id))
|
||||
if cached_rank is not None:
|
||||
return max(0.0, min(cached_rank, 1.0))
|
||||
|
||||
# Hard-priority semantics:
|
||||
# 1) Compare by preset[0] metric first;
|
||||
# 2) only if tied, compare preset[1], preset[2], ...
|
||||
# 3) if all preset metrics tie and LRU is enabled, use LRU as final tiebreak.
|
||||
metric_vectors: dict[str, tuple[float, ...]] = {}
|
||||
for kid in all_key_ids:
|
||||
vector_parts: list[float] = []
|
||||
for preset_name, mode in presets:
|
||||
metric = 0.5
|
||||
dim = get_preset_dimension(preset_name)
|
||||
if dim is not None:
|
||||
metric = dim.compute_metric(
|
||||
key_id=kid,
|
||||
all_key_ids=all_key_ids,
|
||||
keys_by_id=keys_by_id,
|
||||
lru_scores=lru_scores,
|
||||
context=context,
|
||||
mode=mode,
|
||||
)
|
||||
metric_value = safe_float(metric)
|
||||
vector_parts.append(
|
||||
max(0.0, min(metric_value, 1.0)) if metric_value is not None else 0.5
|
||||
)
|
||||
|
||||
if lru_enabled:
|
||||
vector_parts.append(rank_ascending(kid, lru_scores, all_key_ids))
|
||||
|
||||
metric_vectors[kid] = tuple(vector_parts)
|
||||
|
||||
decorated = [
|
||||
(metric_vectors.get(kid, (0.5,)), idx, kid) for idx, kid in enumerate(all_key_ids)
|
||||
]
|
||||
decorated.sort(key=lambda item: (item[0], item[1]))
|
||||
|
||||
total = len(decorated)
|
||||
ranks: dict[str, float] = {}
|
||||
for rank_idx, (_vec, _idx, kid) in enumerate(decorated):
|
||||
ranks[kid] = 0.0 if total <= 1 else rank_idx / float(total - 1)
|
||||
|
||||
context["_preset_hard_order_cache"] = {
|
||||
"signature": cache_signature,
|
||||
"ranks": ranks,
|
||||
}
|
||||
rank = safe_float(ranks.get(key_id))
|
||||
if rank is None:
|
||||
return 0.5
|
||||
return max(0.0, min(rank, 1.0))
|
||||
|
||||
|
||||
register_pool_strategy("multi_score", MultiScoreStrategy())
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MultiScoreStrategy",
|
||||
]
|
||||
107
_deprecated_py_src/services/provider/pool/strategy.py
Normal file
107
_deprecated_py_src/services/provider/pool/strategy.py
Normal file
@@ -0,0 +1,107 @@
|
||||
"""Pluggable pool scheduling strategies.
|
||||
|
||||
Strategies allow customising pool-level candidate selection without
|
||||
modifying the core :class:`PoolManager`. Each strategy is an object
|
||||
that implements one or more optional methods defined by the
|
||||
:class:`PoolSchedulingStrategy` protocol.
|
||||
|
||||
Registration uses a thread-safe global registry (same pattern as
|
||||
:mod:`~src.services.provider.pool.hooks`).
|
||||
|
||||
Usage::
|
||||
|
||||
from src.services.provider.pool.strategy import register_pool_strategy
|
||||
|
||||
class MyStrategy:
|
||||
name = "usage_weight"
|
||||
|
||||
def compute_score(self, *, key_id, config, context):
|
||||
...
|
||||
|
||||
register_pool_strategy("usage_weight", MyStrategy())
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
from src.services.provider.pool.trace import PoolCandidateTrace
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class PoolSchedulingStrategy(Protocol):
|
||||
"""Pluggable pool scheduling strategy.
|
||||
|
||||
All methods are optional -- callers check via ``hasattr``.
|
||||
Strategies are activated per-provider through ``PoolConfig.strategies``.
|
||||
"""
|
||||
|
||||
name: str
|
||||
|
||||
def on_before_select(
|
||||
self,
|
||||
*,
|
||||
provider_id: str,
|
||||
key_ids: list[str],
|
||||
config: PoolConfig,
|
||||
context: dict[str, Any],
|
||||
) -> list[str] | None:
|
||||
"""Filter / reorder *key_ids* before selection.
|
||||
|
||||
Return ``None`` to leave the list unchanged.
|
||||
"""
|
||||
...
|
||||
|
||||
def on_after_select(
|
||||
self,
|
||||
*,
|
||||
provider_id: str,
|
||||
selected_key_id: str,
|
||||
trace: PoolCandidateTrace,
|
||||
config: PoolConfig,
|
||||
context: dict[str, Any],
|
||||
) -> None:
|
||||
"""Called after a key has been selected (for logging / metrics)."""
|
||||
...
|
||||
|
||||
def compute_score(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
config: PoolConfig,
|
||||
context: dict[str, Any],
|
||||
) -> float | None:
|
||||
"""Return a custom sort score. ``None`` means "do not override"."""
|
||||
...
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_strategy_registry: dict[str, PoolSchedulingStrategy] = {}
|
||||
_strategy_lock = threading.Lock()
|
||||
|
||||
|
||||
def register_pool_strategy(name: str, strategy: PoolSchedulingStrategy) -> None:
|
||||
"""Register a pool scheduling strategy globally."""
|
||||
with _strategy_lock:
|
||||
_strategy_registry[name] = strategy
|
||||
|
||||
|
||||
def get_pool_strategy(name: str) -> PoolSchedulingStrategy | None:
|
||||
"""Return a registered strategy by *name*, or ``None``."""
|
||||
return _strategy_registry.get(name)
|
||||
|
||||
|
||||
def get_active_strategies(names: tuple[str, ...] | list[str]) -> list[PoolSchedulingStrategy]:
|
||||
"""Return registered strategies whose names appear in *names*."""
|
||||
result: list[PoolSchedulingStrategy] = []
|
||||
for n in names:
|
||||
s = _strategy_registry.get(n)
|
||||
if s is not None:
|
||||
result.append(s)
|
||||
return result
|
||||
137
_deprecated_py_src/services/provider/pool/trace.py
Normal file
137
_deprecated_py_src/services/provider/pool/trace.py
Normal file
@@ -0,0 +1,137 @@
|
||||
"""Pool scheduling trace -- per-candidate decision records.
|
||||
|
||||
Collects scheduling decisions made during pool-level candidate selection
|
||||
without adding any extra Redis round-trips. Trace data is later written
|
||||
to ``RequestCandidate.extra_data`` and ``Usage.request_metadata``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PoolCandidateTrace:
|
||||
"""Single candidate scheduling decision in a pool context."""
|
||||
|
||||
key_id: str
|
||||
reason: str = "" # sticky / lru / random / tiebreak
|
||||
sticky_hit: bool = False
|
||||
lru_score: float = 0.0
|
||||
cost_window_usage: int = 0
|
||||
cost_limit: int | None = None
|
||||
cost_soft_threshold: bool = False
|
||||
skipped: bool = False
|
||||
skip_type: str | None = None # cooldown / cost_exhausted / account_blocked / upstream
|
||||
cooldown_reason: str | None = None
|
||||
cooldown_ttl: int | None = None
|
||||
account_block_code: str | None = None
|
||||
account_block_label: str | None = None
|
||||
account_block_reason: str | None = None
|
||||
latency_avg_ms: float = 0.0
|
||||
health_score: float = 1.0
|
||||
composite_score: float = 0.0
|
||||
scoring_mode: str = "lru"
|
||||
|
||||
def to_extra_data(self) -> dict[str, Any]:
|
||||
"""Build dict to merge into ``RequestCandidate.extra_data``."""
|
||||
if self.skipped:
|
||||
skip_info: dict[str, Any] = {"type": self.skip_type}
|
||||
if self.cooldown_reason is not None:
|
||||
skip_info["cooldown_reason"] = self.cooldown_reason
|
||||
if self.cooldown_ttl is not None:
|
||||
skip_info["cooldown_ttl"] = self.cooldown_ttl
|
||||
if self.account_block_code is not None:
|
||||
skip_info["account_block_code"] = self.account_block_code
|
||||
if self.account_block_label is not None:
|
||||
skip_info["account_block_label"] = self.account_block_label
|
||||
if self.account_block_reason is not None:
|
||||
skip_info["account_block_reason"] = self.account_block_reason
|
||||
if self.cost_window_usage:
|
||||
skip_info["cost_window_usage"] = self.cost_window_usage
|
||||
if self.scoring_mode:
|
||||
skip_info["scoring_mode"] = self.scoring_mode
|
||||
return {"pool_skip": skip_info}
|
||||
|
||||
sel: dict[str, Any] = {"reason": self.reason}
|
||||
if self.sticky_hit:
|
||||
sel["sticky_hit"] = True
|
||||
if self.lru_score:
|
||||
sel["lru_score"] = self.lru_score
|
||||
if self.cost_window_usage:
|
||||
sel["cost_window_usage"] = self.cost_window_usage
|
||||
if self.cost_limit is not None:
|
||||
sel["cost_limit"] = self.cost_limit
|
||||
if self.cost_soft_threshold:
|
||||
sel["cost_soft_threshold"] = True
|
||||
if self.latency_avg_ms > 0:
|
||||
sel["latency_avg_ms"] = round(self.latency_avg_ms, 2)
|
||||
if self.health_score < 1.0:
|
||||
sel["health_score"] = round(self.health_score, 4)
|
||||
if self.reason == "multi_score":
|
||||
sel["composite_score"] = round(self.composite_score, 6)
|
||||
if self.scoring_mode:
|
||||
sel["scoring_mode"] = self.scoring_mode
|
||||
return {"pool_selection": sel}
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PoolSchedulingTrace:
|
||||
"""Aggregated scheduling trace for one pool-provider dispatch."""
|
||||
|
||||
provider_id: str
|
||||
total_keys: int = 0
|
||||
sticky_session_used: bool = False
|
||||
session_uuid: str | None = None
|
||||
candidate_traces: dict[str, PoolCandidateTrace] = field(default_factory=dict)
|
||||
|
||||
def build_summary(
|
||||
self,
|
||||
success_key_id: str | None = None,
|
||||
*,
|
||||
attempted_key_ids: set[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build compact dict for ``Usage.request_metadata["pool_summary"]``."""
|
||||
skipped_cooldown = 0
|
||||
skipped_cost = 0
|
||||
skipped_account_blocked = 0
|
||||
attempted = 0
|
||||
for t in self.candidate_traces.values():
|
||||
if t.skipped:
|
||||
if t.skip_type == "cooldown":
|
||||
skipped_cooldown += 1
|
||||
elif t.skip_type == "cost_exhausted":
|
||||
skipped_cost += 1
|
||||
elif t.skip_type == "account_blocked":
|
||||
skipped_account_blocked += 1
|
||||
|
||||
if attempted_key_ids is None:
|
||||
# Backward-compatible behavior: count all schedulable keys.
|
||||
attempted = sum(1 for t in self.candidate_traces.values() if not t.skipped)
|
||||
else:
|
||||
# Preferred behavior: count only keys that were actually executed.
|
||||
attempted = sum(
|
||||
1
|
||||
for kid in attempted_key_ids
|
||||
if kid in self.candidate_traces and not self.candidate_traces[kid].skipped
|
||||
)
|
||||
|
||||
success_reason: str | None = None
|
||||
if success_key_id and success_key_id in self.candidate_traces:
|
||||
success_reason = self.candidate_traces[success_key_id].reason
|
||||
|
||||
summary: dict[str, Any] = {
|
||||
"enabled": True,
|
||||
"total_keys": self.total_keys,
|
||||
"attempted": attempted,
|
||||
"skipped_cooldown": skipped_cooldown,
|
||||
"skipped_cost": skipped_cost,
|
||||
"skipped_account_blocked": skipped_account_blocked,
|
||||
"sticky_session": self.sticky_session_used,
|
||||
}
|
||||
if success_key_id:
|
||||
summary["success_key_id"] = success_key_id[:8]
|
||||
if success_reason:
|
||||
summary["success_reason"] = success_reason
|
||||
return summary
|
||||
Reference in New Issue
Block a user