mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
refactor: 修复候选分页逻辑并拆分配额检查模块
- list_all_candidates 返回 provider_batch_count,区分"无候选"与"无 Provider" 避免分页在有 Provider 但无候选时提前终止 - 将配额检查逻辑从 aware_scheduler.py 拆分到 quota_skipper.py - 提取 reorder_candidates 方法,支持跨页汇总后全局重排序 - CandidateResolver 分页循环改用 provider_batch_count 判断终止 - candidate extra_data 中新增 mapping_matched_model 字段 - 新增契约测试和分页行为测试
This commit is contained in:
@@ -216,7 +216,7 @@ async def _select_provider_candidate(
|
|||||||
# 要求 gemini_files 能力:只有 Google 官方 API 才支持 Files API
|
# 要求 gemini_files 能力:只有 Google 官方 API 才支持 Files API
|
||||||
capability_requirements = {"gemini_files": True} if require_files_capability else None
|
capability_requirements = {"gemini_files": True} if require_files_capability else None
|
||||||
|
|
||||||
candidates, _global_model_id = await scheduler.list_all_candidates(
|
candidates, _global_model_id, _provider_count = await scheduler.list_all_candidates(
|
||||||
db=db,
|
db=db,
|
||||||
api_format="gemini:chat",
|
api_format="gemini:chat",
|
||||||
model_name=model_name,
|
model_name=model_name,
|
||||||
|
|||||||
196
src/services/cache/aware_scheduler.py
vendored
196
src/services/cache/aware_scheduler.py
vendored
@@ -54,7 +54,6 @@ from src.core.model_permissions import (
|
|||||||
get_allowed_models_preview,
|
get_allowed_models_preview,
|
||||||
merge_allowed_models,
|
merge_allowed_models,
|
||||||
)
|
)
|
||||||
from src.core.provider_types import ProviderType, normalize_provider_type
|
|
||||||
from src.models.database import (
|
from src.models.database import (
|
||||||
ApiKey,
|
ApiKey,
|
||||||
Model,
|
Model,
|
||||||
@@ -62,6 +61,7 @@ from src.models.database import (
|
|||||||
ProviderAPIKey,
|
ProviderAPIKey,
|
||||||
ProviderEndpoint,
|
ProviderEndpoint,
|
||||||
)
|
)
|
||||||
|
from src.services.cache.quota_skipper import is_key_quota_exhausted
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from src.models.database import GlobalModel
|
from src.models.database import GlobalModel
|
||||||
@@ -134,106 +134,6 @@ class ProviderCandidate:
|
|||||||
return self._stable_order_key() < other._stable_order_key()
|
return self._stable_order_key() < other._stable_order_key()
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Quota-based skipping
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def _pct_is_exhausted(value: object) -> bool:
|
|
||||||
"""Return True when used_percent indicates 0% remaining."""
|
|
||||||
try:
|
|
||||||
pct = float(value) # type: ignore[arg-type]
|
|
||||||
except (TypeError, ValueError):
|
|
||||||
return False
|
|
||||||
# Some upstreams may return values slightly above 100 due to rounding.
|
|
||||||
return pct >= 100.0 - 1e-6
|
|
||||||
|
|
||||||
|
|
||||||
def _float_or_none(value: object) -> float | None:
|
|
||||||
try:
|
|
||||||
if value is None:
|
|
||||||
return None
|
|
||||||
return float(value) # type: ignore[arg-type]
|
|
||||||
except (TypeError, ValueError):
|
|
||||||
return None
|
|
||||||
|
|
||||||
|
|
||||||
def _is_key_quota_exhausted(
|
|
||||||
provider_type: str | None,
|
|
||||||
key: ProviderAPIKey,
|
|
||||||
*,
|
|
||||||
model_name: str,
|
|
||||||
) -> tuple[bool, str | None]:
|
|
||||||
"""Check ProviderAPIKey.upstream_metadata quota and decide whether to skip.
|
|
||||||
|
|
||||||
Requirements:
|
|
||||||
- Kiro: account-level quota. When remaining == 0, skip this key; allow again when remaining > 0.
|
|
||||||
- Codex: only consider weekly quota + 5H quota (ignore code review quota).
|
|
||||||
If either remaining is 0%, skip this key.
|
|
||||||
- Antigravity: quota is per-model; do not disable the account.
|
|
||||||
When the requested model's quota is 0%, skip this key.
|
|
||||||
"""
|
|
||||||
pt = normalize_provider_type(provider_type)
|
|
||||||
|
|
||||||
upstream = getattr(key, "upstream_metadata", None) or {}
|
|
||||||
if not isinstance(upstream, dict):
|
|
||||||
return False, None
|
|
||||||
|
|
||||||
if pt == ProviderType.KIRO:
|
|
||||||
kiro_meta = upstream.get("kiro")
|
|
||||||
if not isinstance(kiro_meta, dict):
|
|
||||||
return False, None
|
|
||||||
|
|
||||||
remaining = _float_or_none(kiro_meta.get("remaining"))
|
|
||||||
|
|
||||||
if remaining is not None and remaining <= 0.0:
|
|
||||||
return True, "Kiro 账号配额剩余 0"
|
|
||||||
|
|
||||||
return False, None
|
|
||||||
|
|
||||||
if pt == ProviderType.CODEX:
|
|
||||||
codex_meta = upstream.get("codex")
|
|
||||||
if not isinstance(codex_meta, dict):
|
|
||||||
return False, None
|
|
||||||
|
|
||||||
weekly_used = codex_meta.get("primary_used_percent")
|
|
||||||
five_hour_used = codex_meta.get("secondary_used_percent")
|
|
||||||
|
|
||||||
exhausted_parts: list[str] = []
|
|
||||||
if _pct_is_exhausted(weekly_used):
|
|
||||||
exhausted_parts.append("周限额剩余 0%")
|
|
||||||
if _pct_is_exhausted(five_hour_used):
|
|
||||||
exhausted_parts.append("5H 限额剩余 0%")
|
|
||||||
|
|
||||||
if exhausted_parts:
|
|
||||||
return True, "Codex " + ",".join(exhausted_parts)
|
|
||||||
|
|
||||||
return False, None
|
|
||||||
|
|
||||||
if pt == ProviderType.ANTIGRAVITY:
|
|
||||||
ag_meta = upstream.get("antigravity")
|
|
||||||
if not isinstance(ag_meta, dict):
|
|
||||||
return False, None
|
|
||||||
|
|
||||||
quota_by_model = ag_meta.get("quota_by_model")
|
|
||||||
if not isinstance(quota_by_model, dict):
|
|
||||||
return False, None
|
|
||||||
|
|
||||||
model_quota = quota_by_model.get(model_name)
|
|
||||||
if not isinstance(model_quota, dict):
|
|
||||||
return False, None
|
|
||||||
|
|
||||||
remaining_fraction = _float_or_none(model_quota.get("remaining_fraction"))
|
|
||||||
if remaining_fraction is not None and remaining_fraction <= 0.0:
|
|
||||||
return True, f"Antigravity 模型 {model_name} 配额剩余 0%"
|
|
||||||
if _pct_is_exhausted(model_quota.get("used_percent")):
|
|
||||||
return True, f"Antigravity 模型 {model_name} 配额剩余 0%"
|
|
||||||
|
|
||||||
return False, None
|
|
||||||
|
|
||||||
return False, None
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class ConcurrencySnapshot:
|
class ConcurrencySnapshot:
|
||||||
key_current: int
|
key_current: int
|
||||||
@@ -440,20 +340,22 @@ class CacheAwareScheduler:
|
|||||||
global_model_id = None # 用于缓存亲和性
|
global_model_id = None # 用于缓存亲和性
|
||||||
|
|
||||||
while True:
|
while True:
|
||||||
candidates, resolved_global_model_id = await self.list_all_candidates(
|
candidates, resolved_global_model_id, provider_batch_count = (
|
||||||
db=db,
|
await self.list_all_candidates(
|
||||||
api_format=normalized_format,
|
db=db,
|
||||||
model_name=model_name,
|
api_format=normalized_format,
|
||||||
affinity_key=affinity_key,
|
model_name=model_name,
|
||||||
provider_offset=provider_offset,
|
affinity_key=affinity_key,
|
||||||
provider_limit=provider_batch_size,
|
provider_offset=provider_offset,
|
||||||
max_candidates=max_candidates_per_batch,
|
provider_limit=provider_batch_size,
|
||||||
|
max_candidates=max_candidates_per_batch,
|
||||||
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
if resolved_global_model_id and global_model_id is None:
|
if resolved_global_model_id and global_model_id is None:
|
||||||
global_model_id = resolved_global_model_id
|
global_model_id = resolved_global_model_id
|
||||||
|
|
||||||
if not candidates:
|
if provider_batch_count == 0:
|
||||||
if provider_offset == 0:
|
if provider_offset == 0:
|
||||||
# 没有找到任何候选,提供友好的错误提示(不暴露内部信息)
|
# 没有找到任何候选,提供友好的错误提示(不暴露内部信息)
|
||||||
raise ProviderNotAvailableException("请求的模型当前不可用")
|
raise ProviderNotAvailableException("请求的模型当前不可用")
|
||||||
@@ -513,6 +415,8 @@ class CacheAwareScheduler:
|
|||||||
return provider, endpoint, key
|
return provider, endpoint, key
|
||||||
|
|
||||||
provider_offset += provider_batch_size
|
provider_offset += provider_batch_size
|
||||||
|
if provider_batch_count < provider_batch_size:
|
||||||
|
break
|
||||||
|
|
||||||
raise ProviderNotAvailableException("服务暂时繁忙,请稍后重试")
|
raise ProviderNotAvailableException("服务暂时繁忙,请稍后重试")
|
||||||
|
|
||||||
@@ -716,7 +620,7 @@ class CacheAwareScheduler:
|
|||||||
max_candidates: int | None = None,
|
max_candidates: int | None = None,
|
||||||
is_stream: bool = False,
|
is_stream: bool = False,
|
||||||
capability_requirements: dict[str, bool] | None = None,
|
capability_requirements: dict[str, bool] | None = None,
|
||||||
) -> tuple[list[ProviderCandidate], str]:
|
) -> tuple[list[ProviderCandidate], str, int]:
|
||||||
"""
|
"""
|
||||||
预先获取所有可用的 Provider/Endpoint/Key 组合
|
预先获取所有可用的 Provider/Endpoint/Key 组合
|
||||||
|
|
||||||
@@ -738,7 +642,9 @@ class CacheAwareScheduler:
|
|||||||
capability_requirements: 能力需求(用于过滤不满足能力要求的 Key)
|
capability_requirements: 能力需求(用于过滤不满足能力要求的 Key)
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
(候选列表, global_model_id) - global_model_id 用于缓存亲和性
|
(候选列表, global_model_id, provider_batch_count)
|
||||||
|
- global_model_id 用于缓存亲和性
|
||||||
|
- provider_batch_count 表示本次查询到的 Provider 数量(未应用 allowed_providers 过滤前)
|
||||||
"""
|
"""
|
||||||
# If the caller already touched the DB, release the connection before we do async work.
|
# If the caller already touched the DB, release the connection before we do async work.
|
||||||
self._release_db_connection_before_await(db)
|
self._release_db_connection_before_await(db)
|
||||||
@@ -772,6 +678,8 @@ class CacheAwareScheduler:
|
|||||||
# 使用 GlobalModel.id 作为缓存亲和性的模型标识,确保映射名和规范名都能命中同一个缓存
|
# 使用 GlobalModel.id 作为缓存亲和性的模型标识,确保映射名和规范名都能命中同一个缓存
|
||||||
global_model_id: str = str(global_model.id)
|
global_model_id: str = str(global_model.id)
|
||||||
|
|
||||||
|
queried_provider_count = 0
|
||||||
|
|
||||||
# 提取模型映射(用于 Provider Key 的 allowed_models 匹配)
|
# 提取模型映射(用于 Provider Key 的 allowed_models 匹配)
|
||||||
model_mappings: list[str] = (global_model.config or {}).get("model_mappings", [])
|
model_mappings: list[str] = (global_model.config or {}).get("model_mappings", [])
|
||||||
if model_mappings:
|
if model_mappings:
|
||||||
@@ -793,7 +701,7 @@ class CacheAwareScheduler:
|
|||||||
f"API Key {user_api_key.id[:8] if user_api_key else 'N/A'}... 不允许使用 API 格式 {target_format}, "
|
f"API Key {user_api_key.id[:8] if user_api_key else 'N/A'}... 不允许使用 API 格式 {target_format}, "
|
||||||
f"允许的格式: {allowed_api_formats}"
|
f"允许的格式: {allowed_api_formats}"
|
||||||
)
|
)
|
||||||
return [], global_model_id
|
return [], global_model_id, queried_provider_count
|
||||||
|
|
||||||
# 0.2 检查模型是否被允许
|
# 0.2 检查模型是否被允许
|
||||||
if not check_model_allowed(
|
if not check_model_allowed(
|
||||||
@@ -804,7 +712,7 @@ class CacheAwareScheduler:
|
|||||||
f"用户/API Key 不允许使用模型 {model_name}, "
|
f"用户/API Key 不允许使用模型 {model_name}, "
|
||||||
f"允许的模型: {get_allowed_models_preview(allowed_models)}"
|
f"允许的模型: {get_allowed_models_preview(allowed_models)}"
|
||||||
)
|
)
|
||||||
return [], global_model_id
|
return [], global_model_id, queried_provider_count
|
||||||
|
|
||||||
# 1. 查询 Providers
|
# 1. 查询 Providers
|
||||||
providers = self._query_providers(
|
providers = self._query_providers(
|
||||||
@@ -812,6 +720,7 @@ class CacheAwareScheduler:
|
|||||||
provider_offset=provider_offset,
|
provider_offset=provider_offset,
|
||||||
provider_limit=provider_limit,
|
provider_limit=provider_limit,
|
||||||
)
|
)
|
||||||
|
queried_provider_count = len(providers)
|
||||||
|
|
||||||
# Provider query starts a transaction; release connection before entering async candidate build.
|
# Provider query starts a transaction; release connection before entering async candidate build.
|
||||||
self._release_db_connection_before_await(db)
|
self._release_db_connection_before_await(db)
|
||||||
@@ -831,7 +740,7 @@ class CacheAwareScheduler:
|
|||||||
)
|
)
|
||||||
|
|
||||||
if not providers:
|
if not providers:
|
||||||
return [], global_model_id
|
return [], global_model_id, queried_provider_count
|
||||||
|
|
||||||
# 1.5 根据 allowed_providers 过滤(合并 ApiKey 和 User 的限制)
|
# 1.5 根据 allowed_providers 过滤(合并 ApiKey 和 User 的限制)
|
||||||
if allowed_providers is not None:
|
if allowed_providers is not None:
|
||||||
@@ -844,7 +753,7 @@ class CacheAwareScheduler:
|
|||||||
logger.debug(f"用户/API Key 过滤 Provider: {original_count} -> {len(providers)}")
|
logger.debug(f"用户/API Key 过滤 Provider: {original_count} -> {len(providers)}")
|
||||||
|
|
||||||
if not providers:
|
if not providers:
|
||||||
return [], global_model_id
|
return [], global_model_id, queried_provider_count
|
||||||
|
|
||||||
# 2. 构建候选列表(传入 is_stream 和 capability_requirements 用于过滤)
|
# 2. 构建候选列表(传入 is_stream 和 capability_requirements 用于过滤)
|
||||||
|
|
||||||
@@ -864,8 +773,14 @@ class CacheAwareScheduler:
|
|||||||
global_conversion_enabled=global_conversion_enabled,
|
global_conversion_enabled=global_conversion_enabled,
|
||||||
)
|
)
|
||||||
|
|
||||||
# 3. 应用优先级模式排序
|
# 3. 应用优先级模式排序 + 调度模式排序
|
||||||
candidates = self._apply_priority_mode_sort(candidates, db, affinity_key, target_format)
|
candidates = await self.reorder_candidates(
|
||||||
|
candidates=candidates,
|
||||||
|
db=db,
|
||||||
|
affinity_key=affinity_key,
|
||||||
|
api_format=target_format,
|
||||||
|
global_model_id=global_model_id,
|
||||||
|
)
|
||||||
|
|
||||||
# 更新指标
|
# 更新指标
|
||||||
self._metrics["total_candidates"] += len(candidates)
|
self._metrics["total_candidates"] += len(candidates)
|
||||||
@@ -876,28 +791,55 @@ class CacheAwareScheduler:
|
|||||||
f"(api_format={target_format}, model={model_name})"
|
f"(api_format={target_format}, model={model_name})"
|
||||||
)
|
)
|
||||||
|
|
||||||
# 4. 根据调度模式应用不同的排序策略
|
return candidates, global_model_id, queried_provider_count
|
||||||
|
|
||||||
|
async def reorder_candidates(
|
||||||
|
self,
|
||||||
|
candidates: list[ProviderCandidate],
|
||||||
|
db: Session,
|
||||||
|
affinity_key: str | None = None,
|
||||||
|
api_format: str | None = None,
|
||||||
|
global_model_id: str | None = None,
|
||||||
|
) -> list[ProviderCandidate]:
|
||||||
|
"""对候选列表应用优先级模式排序和调度模式排序。
|
||||||
|
|
||||||
|
在分页汇总后调用此方法可修正跨页排序失真。
|
||||||
|
|
||||||
|
Args:
|
||||||
|
candidates: 候选列表
|
||||||
|
db: 数据库会话
|
||||||
|
affinity_key: 亲和性标识符
|
||||||
|
api_format: API 格式
|
||||||
|
global_model_id: GlobalModel ID(缓存亲和模式需要)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
重排序后的候选列表
|
||||||
|
"""
|
||||||
|
if not candidates:
|
||||||
|
return candidates
|
||||||
|
|
||||||
|
# 1. 优先级模式排序
|
||||||
|
candidates = self._apply_priority_mode_sort(candidates, db, affinity_key, api_format)
|
||||||
|
|
||||||
|
# 2. 调度模式排序
|
||||||
if self.scheduling_mode == self.SCHEDULING_MODE_CACHE_AFFINITY:
|
if self.scheduling_mode == self.SCHEDULING_MODE_CACHE_AFFINITY:
|
||||||
# 缓存亲和模式:优先使用缓存的,同优先级内哈希分散
|
if affinity_key and candidates and global_model_id:
|
||||||
if affinity_key and candidates:
|
|
||||||
candidates = await self._apply_cache_affinity(
|
candidates = await self._apply_cache_affinity(
|
||||||
candidates=candidates,
|
candidates=candidates,
|
||||||
db=db,
|
db=db,
|
||||||
affinity_key=affinity_key,
|
affinity_key=affinity_key,
|
||||||
api_format=target_format,
|
api_format=api_format or "",
|
||||||
global_model_id=global_model_id,
|
global_model_id=global_model_id,
|
||||||
)
|
)
|
||||||
elif self.scheduling_mode == self.SCHEDULING_MODE_LOAD_BALANCE:
|
elif self.scheduling_mode == self.SCHEDULING_MODE_LOAD_BALANCE:
|
||||||
# 负载均衡模式:忽略缓存,同优先级内随机轮换
|
candidates = self._apply_load_balance(candidates, api_format)
|
||||||
candidates = self._apply_load_balance(candidates, target_format)
|
|
||||||
for candidate in candidates:
|
for candidate in candidates:
|
||||||
candidate.is_cached = False
|
candidate.is_cached = False
|
||||||
else:
|
else:
|
||||||
# 固定顺序模式:严格按优先级,忽略缓存
|
|
||||||
for candidate in candidates:
|
for candidate in candidates:
|
||||||
candidate.is_cached = False
|
candidate.is_cached = False
|
||||||
|
|
||||||
return candidates, global_model_id
|
return candidates
|
||||||
|
|
||||||
def _query_providers(
|
def _query_providers(
|
||||||
self,
|
self,
|
||||||
@@ -1173,7 +1115,7 @@ class CacheAwareScheduler:
|
|||||||
|
|
||||||
effective_model_name = mapping_matched_model or model_name
|
effective_model_name = mapping_matched_model or model_name
|
||||||
|
|
||||||
quota_exhausted, quota_reason = _is_key_quota_exhausted(
|
quota_exhausted, quota_reason = is_key_quota_exhausted(
|
||||||
provider_type,
|
provider_type,
|
||||||
key,
|
key,
|
||||||
model_name=effective_model_name,
|
model_name=effective_model_name,
|
||||||
|
|||||||
100
src/services/cache/quota_skipper.py
vendored
Normal file
100
src/services/cache/quota_skipper.py
vendored
Normal file
@@ -0,0 +1,100 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from src.core.provider_types import ProviderType, normalize_provider_type
|
||||||
|
from src.models.database import ProviderAPIKey
|
||||||
|
|
||||||
|
|
||||||
|
def _pct_is_exhausted(value: object) -> bool:
|
||||||
|
"""Return True when used_percent indicates 0% remaining."""
|
||||||
|
try:
|
||||||
|
pct = float(value) # type: ignore[arg-type]
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return False
|
||||||
|
# Some upstreams may return values slightly above 100 due to rounding.
|
||||||
|
return pct >= 100.0 - 1e-6
|
||||||
|
|
||||||
|
|
||||||
|
def _float_or_none(value: object) -> float | None:
|
||||||
|
try:
|
||||||
|
if value is None:
|
||||||
|
return None
|
||||||
|
return float(value) # type: ignore[arg-type]
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def is_key_quota_exhausted(
|
||||||
|
provider_type: str | None,
|
||||||
|
key: ProviderAPIKey,
|
||||||
|
*,
|
||||||
|
model_name: str,
|
||||||
|
) -> tuple[bool, str | None]:
|
||||||
|
"""Check ProviderAPIKey.upstream_metadata quota and decide whether to skip.
|
||||||
|
|
||||||
|
Requirements:
|
||||||
|
- Kiro: account-level quota. When remaining == 0, skip this key; allow again when remaining > 0.
|
||||||
|
- Codex: only consider weekly quota + 5H quota (ignore code review quota).
|
||||||
|
If either remaining is 0%, skip this key.
|
||||||
|
- Antigravity: quota is per-model; do not disable the account.
|
||||||
|
When the requested model's quota is 0%, skip this key.
|
||||||
|
"""
|
||||||
|
|
||||||
|
pt = normalize_provider_type(provider_type)
|
||||||
|
|
||||||
|
upstream = getattr(key, "upstream_metadata", None) or {}
|
||||||
|
if not isinstance(upstream, dict):
|
||||||
|
return False, None
|
||||||
|
|
||||||
|
if pt == ProviderType.KIRO:
|
||||||
|
kiro_meta = upstream.get("kiro")
|
||||||
|
if not isinstance(kiro_meta, dict):
|
||||||
|
return False, None
|
||||||
|
|
||||||
|
remaining = _float_or_none(kiro_meta.get("remaining"))
|
||||||
|
|
||||||
|
if remaining is not None and remaining <= 0.0:
|
||||||
|
return True, "Kiro 账号配额剩余 0"
|
||||||
|
|
||||||
|
return False, None
|
||||||
|
|
||||||
|
if pt == ProviderType.CODEX:
|
||||||
|
codex_meta = upstream.get("codex")
|
||||||
|
if not isinstance(codex_meta, dict):
|
||||||
|
return False, None
|
||||||
|
|
||||||
|
weekly_used = codex_meta.get("primary_used_percent")
|
||||||
|
five_hour_used = codex_meta.get("secondary_used_percent")
|
||||||
|
|
||||||
|
exhausted_parts: list[str] = []
|
||||||
|
if _pct_is_exhausted(weekly_used):
|
||||||
|
exhausted_parts.append("周限额剩余 0%")
|
||||||
|
if _pct_is_exhausted(five_hour_used):
|
||||||
|
exhausted_parts.append("5H 限额剩余 0%")
|
||||||
|
|
||||||
|
if exhausted_parts:
|
||||||
|
return True, "Codex " + ",".join(exhausted_parts)
|
||||||
|
|
||||||
|
return False, None
|
||||||
|
|
||||||
|
if pt == ProviderType.ANTIGRAVITY:
|
||||||
|
ag_meta = upstream.get("antigravity")
|
||||||
|
if not isinstance(ag_meta, dict):
|
||||||
|
return False, None
|
||||||
|
|
||||||
|
quota_by_model = ag_meta.get("quota_by_model")
|
||||||
|
if not isinstance(quota_by_model, dict):
|
||||||
|
return False, None
|
||||||
|
|
||||||
|
model_quota = quota_by_model.get(model_name)
|
||||||
|
if not isinstance(model_quota, dict):
|
||||||
|
return False, None
|
||||||
|
|
||||||
|
remaining_fraction = _float_or_none(model_quota.get("remaining_fraction"))
|
||||||
|
if remaining_fraction is not None and remaining_fraction <= 0.0:
|
||||||
|
return True, f"Antigravity 模型 {model_name} 配额剩余 0%"
|
||||||
|
if _pct_is_exhausted(model_quota.get("used_percent")):
|
||||||
|
return True, f"Antigravity 模型 {model_name} 配额剩余 0%"
|
||||||
|
|
||||||
|
return False, None
|
||||||
|
|
||||||
|
return False, None
|
||||||
@@ -14,6 +14,7 @@ from src.core.exceptions import ProviderNotAvailableException
|
|||||||
from src.core.logger import logger
|
from src.core.logger import logger
|
||||||
from src.models.database import ApiKey
|
from src.models.database import ApiKey
|
||||||
from src.services.cache.aware_scheduler import CacheAwareScheduler, ProviderCandidate
|
from src.services.cache.aware_scheduler import CacheAwareScheduler, ProviderCandidate
|
||||||
|
from src.services.provider.format import normalize_endpoint_signature
|
||||||
|
|
||||||
|
|
||||||
class CandidateResolver:
|
class CandidateResolver:
|
||||||
@@ -75,45 +76,52 @@ class CandidateResolver:
|
|||||||
provider_offset = 0
|
provider_offset = 0
|
||||||
provider_batch_size = 20
|
provider_batch_size = 20
|
||||||
global_model_id: str | None = None
|
global_model_id: str | None = None
|
||||||
|
api_format_norm = normalize_endpoint_signature(api_format)
|
||||||
|
|
||||||
logger.debug(
|
logger.debug(
|
||||||
"[CandidateResolver] fetch_candidates starting: model={}, api_format={}",
|
"[CandidateResolver] fetch_candidates starting: model={}, api_format={}",
|
||||||
model_name,
|
model_name,
|
||||||
api_format,
|
api_format_norm,
|
||||||
)
|
)
|
||||||
|
|
||||||
while True:
|
while True:
|
||||||
candidates, resolved_global_model_id = await self.cache_scheduler.list_all_candidates(
|
candidates, resolved_global_model_id, provider_batch_count = (
|
||||||
db=self.db,
|
await self.cache_scheduler.list_all_candidates(
|
||||||
api_format=api_format,
|
db=self.db,
|
||||||
model_name=model_name,
|
api_format=api_format_norm,
|
||||||
affinity_key=affinity_key,
|
model_name=model_name,
|
||||||
user_api_key=user_api_key,
|
affinity_key=affinity_key,
|
||||||
provider_offset=provider_offset,
|
user_api_key=user_api_key,
|
||||||
provider_limit=provider_batch_size,
|
provider_offset=provider_offset,
|
||||||
is_stream=is_stream,
|
provider_limit=provider_batch_size,
|
||||||
capability_requirements=capability_requirements,
|
is_stream=is_stream,
|
||||||
|
capability_requirements=capability_requirements,
|
||||||
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
logger.debug(
|
logger.debug(
|
||||||
"[CandidateResolver] list_all_candidates batch: offset={}, returned={} candidates",
|
"[CandidateResolver] list_all_candidates batch: offset={}, providers={}, returned={} candidates",
|
||||||
provider_offset,
|
provider_offset,
|
||||||
|
provider_batch_count,
|
||||||
len(candidates),
|
len(candidates),
|
||||||
)
|
)
|
||||||
|
|
||||||
if resolved_global_model_id and global_model_id is None:
|
if resolved_global_model_id and global_model_id is None:
|
||||||
global_model_id = resolved_global_model_id
|
global_model_id = resolved_global_model_id
|
||||||
|
|
||||||
if not candidates:
|
if provider_batch_count == 0:
|
||||||
break
|
break
|
||||||
|
|
||||||
all_candidates.extend(candidates)
|
all_candidates.extend(candidates)
|
||||||
provider_offset += provider_batch_size
|
provider_offset += provider_batch_size
|
||||||
|
|
||||||
logger.debug(
|
if provider_batch_count < provider_batch_size:
|
||||||
"[CandidateResolver] fetch_candidates completed: total={} candidates",
|
break
|
||||||
len(all_candidates),
|
|
||||||
)
|
logger.debug(
|
||||||
|
"[CandidateResolver] fetch_candidates completed: total={} candidates",
|
||||||
|
len(all_candidates),
|
||||||
|
)
|
||||||
|
|
||||||
if not all_candidates:
|
if not all_candidates:
|
||||||
logger.error(f" [{request_id}] 没有找到任何可用的 Provider/Endpoint/Key 组合")
|
logger.error(f" [{request_id}] 没有找到任何可用的 Provider/Endpoint/Key 组合")
|
||||||
@@ -124,6 +132,22 @@ class CandidateResolver:
|
|||||||
|
|
||||||
logger.debug(f" [{request_id}] 获取到 {len(all_candidates)} 个候选组合")
|
logger.debug(f" [{request_id}] 获取到 {len(all_candidates)} 个候选组合")
|
||||||
|
|
||||||
|
# Provider 分页会导致候选在全局维度上排序失真(尤其是 global_key / 降级分组 / cache_affinity)。
|
||||||
|
# 这里在汇总后再次应用全局排序规则,保证遍历顺序符合当前调度配置。
|
||||||
|
try:
|
||||||
|
all_candidates = await self.cache_scheduler.reorder_candidates(
|
||||||
|
candidates=all_candidates,
|
||||||
|
db=self.db,
|
||||||
|
affinity_key=affinity_key,
|
||||||
|
api_format=api_format_norm,
|
||||||
|
global_model_id=global_model_id,
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
logger.warning(
|
||||||
|
"[CandidateResolver] global reorder failed, keep paged order: {}",
|
||||||
|
exc,
|
||||||
|
)
|
||||||
|
|
||||||
if preferred_key_ids:
|
if preferred_key_ids:
|
||||||
preferred_set = {str(kid) for kid in preferred_key_ids if kid}
|
preferred_set = {str(kid) for kid in preferred_key_ids if kid}
|
||||||
if preferred_set:
|
if preferred_set:
|
||||||
@@ -208,6 +232,7 @@ class CandidateResolver:
|
|||||||
"extra_data": {
|
"extra_data": {
|
||||||
"needs_conversion": candidate.needs_conversion,
|
"needs_conversion": candidate.needs_conversion,
|
||||||
"provider_api_format": candidate.provider_api_format or None,
|
"provider_api_format": candidate.provider_api_format or None,
|
||||||
|
"mapping_matched_model": candidate.mapping_matched_model or None,
|
||||||
},
|
},
|
||||||
"required_capabilities": active_capabilities,
|
"required_capabilities": active_capabilities,
|
||||||
"created_at": datetime.now(timezone.utc),
|
"created_at": datetime.now(timezone.utc),
|
||||||
@@ -241,6 +266,7 @@ class CandidateResolver:
|
|||||||
"extra_data": {
|
"extra_data": {
|
||||||
"needs_conversion": candidate.needs_conversion,
|
"needs_conversion": candidate.needs_conversion,
|
||||||
"provider_api_format": candidate.provider_api_format or None,
|
"provider_api_format": candidate.provider_api_format or None,
|
||||||
|
"mapping_matched_model": candidate.mapping_matched_model or None,
|
||||||
},
|
},
|
||||||
"required_capabilities": active_capabilities,
|
"required_capabilities": active_capabilities,
|
||||||
"created_at": datetime.now(timezone.utc),
|
"created_at": datetime.now(timezone.utc),
|
||||||
|
|||||||
40
tests/contracts/test_public_api_routes_contract.py
Normal file
40
tests/contracts/test_public_api_routes_contract.py
Normal file
@@ -0,0 +1,40 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from fastapi import FastAPI
|
||||||
|
|
||||||
|
|
||||||
|
def _build_contract_app() -> FastAPI:
|
||||||
|
from src.api.public.claude import router as claude_router
|
||||||
|
from src.api.public.gemini import router as gemini_router
|
||||||
|
from src.api.public.openai import router as openai_router
|
||||||
|
|
||||||
|
app = FastAPI()
|
||||||
|
app.include_router(claude_router)
|
||||||
|
app.include_router(openai_router)
|
||||||
|
app.include_router(gemini_router)
|
||||||
|
return app
|
||||||
|
|
||||||
|
|
||||||
|
def test_public_api_routes_contract_paths_and_tags() -> None:
|
||||||
|
app = _build_contract_app()
|
||||||
|
schema = app.openapi()
|
||||||
|
|
||||||
|
paths = schema.get("paths") or {}
|
||||||
|
|
||||||
|
expected = [
|
||||||
|
("/v1/messages", "post", "Claude API"),
|
||||||
|
("/v1/messages/count_tokens", "post", "Claude API"),
|
||||||
|
("/v1/chat/completions", "post", "OpenAI API"),
|
||||||
|
("/v1/responses", "post", "OpenAI API"),
|
||||||
|
("/v1beta/models/{model}:generateContent", "post", "Gemini API"),
|
||||||
|
("/v1beta/models/{model}:streamGenerateContent", "post", "Gemini API"),
|
||||||
|
("/v1/models/{model}:generateContent", "post", "Gemini API"),
|
||||||
|
("/v1/models/{model}:streamGenerateContent", "post", "Gemini API"),
|
||||||
|
]
|
||||||
|
|
||||||
|
for path, method, expected_tag in expected:
|
||||||
|
assert path in paths, f"missing path {path}"
|
||||||
|
operations = paths.get(path) or {}
|
||||||
|
assert method in operations, f"missing {method.upper()} {path}"
|
||||||
|
tags = operations.get(method, {}).get("tags") or []
|
||||||
|
assert expected_tag in tags, f"{method.upper()} {path} missing tag {expected_tag}"
|
||||||
132
tests/contracts/test_scheduler_list_all_candidates_contract.py
Normal file
132
tests/contracts/test_scheduler_list_all_candidates_contract.py
Normal file
@@ -0,0 +1,132 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from types import SimpleNamespace
|
||||||
|
from unittest.mock import AsyncMock, MagicMock, patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from src.services.cache.aware_scheduler import CacheAwareScheduler
|
||||||
|
|
||||||
|
|
||||||
|
def _make_db() -> MagicMock:
|
||||||
|
db = MagicMock()
|
||||||
|
db.new = []
|
||||||
|
db.dirty = []
|
||||||
|
db.deleted = []
|
||||||
|
db.in_transaction.return_value = False
|
||||||
|
return db
|
||||||
|
|
||||||
|
|
||||||
|
def _make_global_model(*, gid: str, name: str) -> SimpleNamespace:
|
||||||
|
return SimpleNamespace(
|
||||||
|
id=gid,
|
||||||
|
name=name,
|
||||||
|
is_active=True,
|
||||||
|
config={},
|
||||||
|
supported_capabilities=[],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_list_all_candidates_returns_provider_batch_count_even_when_candidates_empty() -> (
|
||||||
|
None
|
||||||
|
):
|
||||||
|
"""契约:候选为空不代表 Provider 页为空(用于分页继续拉取下一页)。"""
|
||||||
|
|
||||||
|
scheduler = CacheAwareScheduler()
|
||||||
|
scheduler.scheduling_mode = CacheAwareScheduler.SCHEDULING_MODE_FIXED_ORDER
|
||||||
|
|
||||||
|
db = _make_db()
|
||||||
|
|
||||||
|
providers = [
|
||||||
|
SimpleNamespace(
|
||||||
|
id="p1",
|
||||||
|
name="p1",
|
||||||
|
is_active=True,
|
||||||
|
endpoints=[],
|
||||||
|
models=[],
|
||||||
|
provider_priority=1,
|
||||||
|
),
|
||||||
|
SimpleNamespace(
|
||||||
|
id="p2",
|
||||||
|
name="p2",
|
||||||
|
is_active=True,
|
||||||
|
endpoints=[],
|
||||||
|
models=[],
|
||||||
|
provider_priority=2,
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
# allowed_providers 会把本页的 provider 全过滤掉,导致 candidates 为空;但 provider_batch_count 应保留过滤前数量。
|
||||||
|
user_api_key = SimpleNamespace(
|
||||||
|
id="ak1",
|
||||||
|
allowed_providers=["not-matching"],
|
||||||
|
allowed_models=None,
|
||||||
|
allowed_api_formats=None,
|
||||||
|
user=None,
|
||||||
|
)
|
||||||
|
|
||||||
|
global_model = _make_global_model(gid="gm1", name="gpt-4o")
|
||||||
|
|
||||||
|
with patch.object(scheduler, "_ensure_initialized", new=AsyncMock(return_value=None)):
|
||||||
|
with patch.object(scheduler, "_query_providers", return_value=providers):
|
||||||
|
with patch(
|
||||||
|
"src.services.cache.aware_scheduler.ModelCacheService.get_global_model_by_name",
|
||||||
|
new=AsyncMock(return_value=global_model),
|
||||||
|
):
|
||||||
|
with patch(
|
||||||
|
"src.services.cache.aware_scheduler.SystemConfigService.is_format_conversion_enabled",
|
||||||
|
return_value=True,
|
||||||
|
):
|
||||||
|
candidates, global_model_id, provider_batch_count = (
|
||||||
|
await scheduler.list_all_candidates(
|
||||||
|
db=db,
|
||||||
|
api_format="openai:chat",
|
||||||
|
model_name="gpt-4o",
|
||||||
|
affinity_key=None,
|
||||||
|
user_api_key=user_api_key, # type: ignore[arg-type]
|
||||||
|
provider_offset=0,
|
||||||
|
provider_limit=20,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
assert candidates == []
|
||||||
|
assert global_model_id == "gm1"
|
||||||
|
assert provider_batch_count == 2
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_list_all_candidates_returns_zero_provider_batch_count_when_provider_page_empty() -> (
|
||||||
|
None
|
||||||
|
):
|
||||||
|
scheduler = CacheAwareScheduler()
|
||||||
|
scheduler.scheduling_mode = CacheAwareScheduler.SCHEDULING_MODE_FIXED_ORDER
|
||||||
|
|
||||||
|
db = _make_db()
|
||||||
|
global_model = _make_global_model(gid="gm1", name="gpt-4o")
|
||||||
|
|
||||||
|
with patch.object(scheduler, "_ensure_initialized", new=AsyncMock(return_value=None)):
|
||||||
|
with patch.object(scheduler, "_query_providers", return_value=[]):
|
||||||
|
with patch(
|
||||||
|
"src.services.cache.aware_scheduler.ModelCacheService.get_global_model_by_name",
|
||||||
|
new=AsyncMock(return_value=global_model),
|
||||||
|
):
|
||||||
|
with patch(
|
||||||
|
"src.services.cache.aware_scheduler.SystemConfigService.is_format_conversion_enabled",
|
||||||
|
return_value=True,
|
||||||
|
):
|
||||||
|
candidates, global_model_id, provider_batch_count = (
|
||||||
|
await scheduler.list_all_candidates(
|
||||||
|
db=db,
|
||||||
|
api_format="openai:chat",
|
||||||
|
model_name="gpt-4o",
|
||||||
|
affinity_key=None,
|
||||||
|
user_api_key=None,
|
||||||
|
provider_offset=0,
|
||||||
|
provider_limit=20,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
assert candidates == []
|
||||||
|
assert global_model_id == "gm1"
|
||||||
|
assert provider_batch_count == 0
|
||||||
97
tests/services/test_candidate_resolver_pagination.py
Normal file
97
tests/services/test_candidate_resolver_pagination.py
Normal file
@@ -0,0 +1,97 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from types import SimpleNamespace
|
||||||
|
from typing import Any
|
||||||
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from src.services.cache.aware_scheduler import CacheAwareScheduler
|
||||||
|
from src.services.orchestration.candidate_resolver import CandidateResolver
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeScheduler:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.calls: list[int] = []
|
||||||
|
self.scheduling_mode = CacheAwareScheduler.SCHEDULING_MODE_FIXED_ORDER
|
||||||
|
self.priority_mode = CacheAwareScheduler.PRIORITY_MODE_PROVIDER
|
||||||
|
|
||||||
|
async def list_all_candidates(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
db: Any,
|
||||||
|
api_format: str,
|
||||||
|
model_name: str,
|
||||||
|
affinity_key: str | None = None,
|
||||||
|
user_api_key: Any | None = None,
|
||||||
|
provider_offset: int = 0,
|
||||||
|
provider_limit: int | None = None,
|
||||||
|
max_candidates: int | None = None,
|
||||||
|
is_stream: bool = False,
|
||||||
|
capability_requirements: dict[str, bool] | None = None,
|
||||||
|
) -> tuple[list[Any], str, int]:
|
||||||
|
_ = (
|
||||||
|
db,
|
||||||
|
api_format,
|
||||||
|
model_name,
|
||||||
|
affinity_key,
|
||||||
|
user_api_key,
|
||||||
|
max_candidates,
|
||||||
|
is_stream,
|
||||||
|
capability_requirements,
|
||||||
|
)
|
||||||
|
assert provider_limit is not None
|
||||||
|
|
||||||
|
self.calls.append(int(provider_offset))
|
||||||
|
|
||||||
|
if provider_offset == 0:
|
||||||
|
# Simulate a provider page that has providers but no eligible candidates.
|
||||||
|
return [], "gm1", int(provider_limit)
|
||||||
|
|
||||||
|
if provider_offset == int(provider_limit):
|
||||||
|
# Next page yields one eligible candidate and is also the last provider page.
|
||||||
|
cand = SimpleNamespace(
|
||||||
|
provider=SimpleNamespace(id="p1", name="prov"),
|
||||||
|
endpoint=SimpleNamespace(id="e1"),
|
||||||
|
key=SimpleNamespace(id="k1"),
|
||||||
|
is_skipped=False,
|
||||||
|
skip_reason=None,
|
||||||
|
is_cached=False,
|
||||||
|
needs_conversion=False,
|
||||||
|
provider_api_format=str(api_format),
|
||||||
|
mapping_matched_model=None,
|
||||||
|
)
|
||||||
|
return [cand], "gm1", 5
|
||||||
|
|
||||||
|
return [], "gm1", 0
|
||||||
|
|
||||||
|
async def reorder_candidates(
|
||||||
|
self,
|
||||||
|
candidates: list[Any],
|
||||||
|
db: Any = None,
|
||||||
|
affinity_key: str | None = None,
|
||||||
|
api_format: str | None = None,
|
||||||
|
global_model_id: str | None = None,
|
||||||
|
) -> list[Any]:
|
||||||
|
return candidates
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_candidate_resolver_pagination_continues_on_empty_candidate_batch() -> None:
|
||||||
|
db = MagicMock()
|
||||||
|
scheduler = _FakeScheduler()
|
||||||
|
resolver = CandidateResolver(db=db, cache_scheduler=scheduler) # type: ignore[arg-type]
|
||||||
|
|
||||||
|
candidates, global_model_id = await resolver.fetch_candidates(
|
||||||
|
api_format="openai:chat",
|
||||||
|
model_name="gpt-4o",
|
||||||
|
affinity_key="a1",
|
||||||
|
user_api_key=None,
|
||||||
|
request_id="r1",
|
||||||
|
is_stream=False,
|
||||||
|
capability_requirements=None,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert global_model_id == "gm1"
|
||||||
|
assert len(candidates) == 1
|
||||||
|
assert scheduler.calls == [0, 20]
|
||||||
Reference in New Issue
Block a user