refactor: 移除 Python 后端源码,全面迁移至 Rust gateway 架构

- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本,归档至 _deprecated_py_src/
- 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块,新增 contracts/adaptation 层
- 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构
- 新增 executor 模块,引入 Rust 原生数据库迁移 (aether-data/migrations)
- 简化 CI/Docker 构建流程,移除 base image 二级构建,统一为单一 app image
- 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
This commit is contained in:
fawney19
2026-04-03 16:26:16 +08:00
parent 8f26e1a31f
commit 1d9c77522a
868 changed files with 1735 additions and 2433 deletions

View File

@@ -0,0 +1,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",
]

View 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",
]

View File

@@ -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())

View File

@@ -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())

View File

@@ -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_typeFree 账号优先调度)"
@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())

View File

@@ -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_typeoauth_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())

View File

@@ -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())

View File

@@ -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())

View File

@@ -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())

View File

@@ -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_typePlus/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())

View File

@@ -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())

View File

@@ -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())

View File

@@ -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())

View 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",
]

View File

@@ -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())

View File

@@ -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_typeTeam 账号优先调度)"
@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())