Unify candidate ranking pipeline

This commit is contained in:
fawney19
2026-04-27 17:24:31 +08:00
parent 9b866a6d17
commit 3b542434a2
46 changed files with 3635 additions and 1914 deletions
@@ -144,6 +144,54 @@ async fn read_minimal_candidate_selection_resolves_provider_model_alias() {
assert_eq!(selection[0].selected_provider_model_name, "gpt-5.2");
}
#[tokio::test]
async fn read_minimal_candidate_selection_keeps_all_rows_supporting_requested_model() {
let mut exact = sample_row();
exact.provider_id = "provider-exact".to_string();
exact.endpoint_id = "endpoint-exact".to_string();
exact.key_id = "key-exact".to_string();
exact.model_id = "model-exact".to_string();
exact.global_model_id = "global-exact".to_string();
exact.global_model_name = "gpt-5".to_string();
exact.model_provider_model_name = "gpt-5".to_string();
exact.model_provider_model_mappings = None;
let mut mapped = sample_row();
mapped.provider_id = "provider-mapped".to_string();
mapped.endpoint_id = "endpoint-mapped".to_string();
mapped.key_id = "key-mapped".to_string();
mapped.model_id = "model-mapped".to_string();
mapped.global_model_id = "global-mapped".to_string();
mapped.global_model_name = "claude-sonnet".to_string();
mapped.global_model_mappings = Some(vec!["gpt-5".to_string()]);
mapped.model_provider_model_name = "claude-sonnet-upstream".to_string();
mapped.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "claude-sonnet-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]);
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
exact, mapped,
]));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let state = GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas);
let selection = read_minimal_candidate_selection(&state, "openai:chat", "gpt-5", false, None)
.await
.expect("selection should succeed");
let provider_ids = selection
.iter()
.map(|candidate| candidate.provider_id.as_str())
.collect::<Vec<_>>();
assert_eq!(provider_ids, vec!["provider-exact", "provider-mapped"]);
assert_eq!(
selection[1].selected_provider_model_name,
"claude-sonnet-upstream"
);
}
#[tokio::test]
async fn read_minimal_candidate_selection_allows_resolved_global_model_in_auth_snapshot() {
let mut row = sample_row();
@@ -384,9 +384,9 @@ async fn fixed_order_disables_same_priority_affinity_hash_tiebreaker() {
second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string();
second.key_name = "beta".to_string();
second.provider_priority = 0;
second.provider_priority = 1;
second.key_internal_priority = 0;
second.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
second.key_global_priority_by_format = Some(json!({"openai:chat": 1}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
@@ -438,9 +438,9 @@ async fn cache_affinity_promotes_cached_scheduler_affinity_candidate_when_enable
second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string();
second.key_name = "beta".to_string();
second.provider_priority = 1;
second.provider_priority = 0;
second.key_internal_priority = 0;
second.key_global_priority_by_format = Some(json!({"openai:chat": 1}));
second.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,