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:
fawney19
2026-02-13 13:28:40 +08:00
parent 4c83780b5e
commit d690dcadc7
7 changed files with 482 additions and 145 deletions

View File

@@ -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,

View File

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

View File

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

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

View 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

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