i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-04 08:27:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(provider): support multi-key selection in model tests

Browse Source
This commit is contained in:
zhefox
2026-05-26 13:08:35 +08:00
parent 5fc6dc8019
commit 0bf63cc80e
7 changed files with 253 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs
+25 -24
View File
@@ -1,5 +1,5 @@
use super::super::payload::{
provider_query_extract_api_key_id, provider_query_extract_force_refresh,
provider_query_extract_api_key_ids, provider_query_extract_force_refresh,
provider_query_extract_model, provider_query_extract_provider_id,
provider_query_extract_request_id,
};
@@ -860,7 +860,7 @@ async fn provider_query_select_preferred_non_kiro_endpoint(
provider: &StoredProviderCatalogProvider,
endpoints: &[StoredProviderCatalogEndpoint],
keys: &[StoredProviderCatalogKey],
selected_key_id: Option<&str>,
selected_key_ids: Option<&BTreeSet<String>>,
) -> Option<StoredProviderCatalogEndpoint> {
for priority in 0..=2 {
for endpoint in endpoints.iter().filter(|endpoint| endpoint.is_active) {
@@ -873,7 +873,7 @@ async fn provider_query_select_preferred_non_kiro_endpoint(
}
for key in keys {
if !key.is_active
|| selected_key_id.is_some_and(|value| value != key.id.as_str())
|| !provider_query_selected_key_ids_allow_key(selected_key_ids, &key.id)
|| !provider_query_key_supports_endpoint(
key,
&provider.provider_type,
@@ -905,7 +905,7 @@ async fn provider_query_select_preferred_non_kiro_endpoint(
endpoint.is_active
&& keys.iter().any(|key| {
key.is_active
&& selected_key_id.is_none_or(|value| value == key.id.as_str())
&& provider_query_selected_key_ids_allow_key(selected_key_ids, &key.id)
&& provider_query_key_supports_endpoint(
key,
&provider.provider_type,
@@ -917,6 +917,22 @@ async fn provider_query_select_preferred_non_kiro_endpoint(
.cloned()
}
fn provider_query_selected_key_ids_allow_key(
selected_key_ids: Option<&BTreeSet<String>>,
key_id: &str,
) -> bool {
selected_key_ids.is_none_or(|ids| ids.contains(key_id))
}
fn provider_query_selected_key_ids_all_exist(
selected_key_ids: &BTreeSet<String>,
keys: &[StoredProviderCatalogKey],
) -> bool {
selected_key_ids
.iter()
.all(|id| keys.iter().any(|key| key.id == *id))
}
fn provider_query_test_key_sort_key(
provider_type: &str,
key: &StoredProviderCatalogKey,
@@ -1262,7 +1278,7 @@ async fn provider_query_build_kiro_test_candidates(
ADMIN_PROVIDER_QUERY_NO_ACTIVE_API_KEY_DETAIL,
)
})?;
let selected_key_id = provider_query_extract_api_key_id(payload);
let selected_key_ids = provider_query_extract_api_key_ids(payload);
let requested_endpoint_id = provider_query_extract_endpoint_id(payload);
let requested_api_format = provider_query_extract_api_format(payload);
let endpoint = if requested_endpoint_id.is_none()
@@ -1274,7 +1290,7 @@ async fn provider_query_build_kiro_test_candidates(
provider,
&endpoints,
&all_keys,
selected_key_id.as_deref(),
selected_key_ids.as_ref(),
)
.await
.ok_or_else(|| {
@@ -1307,22 +1323,11 @@ async fn provider_query_build_kiro_test_candidates(
}
};
if let Some(api_key_id) = selected_key_id.as_deref() {
let Some(key) = all_keys.iter().find(|key| key.id == api_key_id) else {
if let Some(selected_key_ids) = selected_key_ids.as_ref() {
if !provider_query_selected_key_ids_all_exist(selected_key_ids, &all_keys) {
return Err(build_admin_provider_query_not_found_response(
ADMIN_PROVIDER_QUERY_API_KEY_NOT_FOUND_DETAIL,
));
};
if !key.is_active
|| !provider_query_key_supports_endpoint(
key,
&provider.provider_type,
&endpoint.api_format,
)
{
return Err(build_admin_provider_query_not_found_response(
ADMIN_PROVIDER_QUERY_NO_ACTIVE_TEST_CANDIDATE_DETAIL,
));
}
}
@@ -1380,11 +1385,7 @@ async fn provider_query_build_kiro_test_candidates(
let mut keys = all_keys
.into_iter()
.filter(|key| key.is_active)
.filter(|key| {
selected_key_id
.as_deref()
.is_none_or(|value| value == key.id.as_str())
})
.filter(|key| provider_query_selected_key_ids_allow_key(selected_key_ids.as_ref(), &key.id))
.filter(|key| {
provider_query_key_supports_endpoint(key, &provider.provider_type, &endpoint.api_format)
})
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs
+21
View File
@@ -231,6 +231,27 @@ fn provider_query_request_body_model_uses_non_empty_string_only() {
);
}
#[test]
fn provider_query_model_test_extracts_multiple_selected_key_ids() {
let payload = json!({
"api_key_ids": [" key-b ", "", "key-a", "key-b"],
"api_key_id": "key-c"
});
let ids = provider_query_extract_api_key_ids(&payload)
.expect("non-empty key selection should be extracted")
.into_iter()
.collect::<Vec<_>>();
assert_eq!(ids, vec!["key-a", "key-b", "key-c"]);
}
#[test]
fn provider_query_model_test_empty_selected_key_ids_keep_default_selection() {
assert!(provider_query_extract_api_key_ids(&json!({})).is_none());
assert!(provider_query_extract_api_key_ids(&json!({ "api_key_ids": [] })).is_none());
}
#[test]
fn provider_query_standard_test_resolves_codex_responses_upstream_streaming() {
assert!(provider_query_resolve_standard_test_upstream_is_stream(
apps/aether-gateway/src/handlers/admin/provider/query/payload.rs
+42
View File
@@ -1,6 +1,7 @@
use axum::body::Bytes;
use axum::response::{IntoResponse, Response};
use serde_json::json;
use std::collections::BTreeSet;
pub(crate) fn parse_admin_provider_query_body(
request_body: Option<&Bytes>,
@@ -36,6 +37,47 @@ pub(crate) fn provider_query_extract_api_key_id(payload: &serde_json::Value) ->
.map(ToOwned::to_owned)
}
fn provider_query_insert_api_key_id(ids: &mut BTreeSet<String>, value: &str) {
let value = value.trim();
if !value.is_empty() {
ids.insert(value.to_string());
}
}
pub(crate) fn provider_query_extract_api_key_ids(
payload: &serde_json::Value,
) -> Option<BTreeSet<String>> {
let mut ids = BTreeSet::new();
if let Some(value) = payload
.get("api_key_ids")
.or_else(|| payload.get("provider_key_ids"))
.or_else(|| payload.get("key_ids"))
{
match value {
serde_json::Value::Array(items) => {
for item in items {
if let Some(value) = item.as_str() {
provider_query_insert_api_key_id(&mut ids, value);
}
}
}
serde_json::Value::String(value) => {
for item in value.split(',') {
provider_query_insert_api_key_id(&mut ids, item);
}
}
_ => {}
}
}
if let Some(api_key_id) = provider_query_extract_api_key_id(payload) {
ids.insert(api_key_id);
}
(!ids.is_empty()).then_some(ids)
}
pub(crate) fn provider_query_extract_force_refresh(payload: &serde_json::Value) -> bool {
payload
.get("force_refresh")