mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-10 13:10:21 +08:00
fix(provider): 模型测试按 Key 模型权限过滤
测试模型前检查 provider key 的 allowed_models: - 空权限视为允许所有模型 - 非空权限需匹配请求模型或映射后的实际模型 - 不匹配的 key 标记为跳过,避免发起测试请求 同时补充相关单测和前端跳过原因文案。
This commit is contained in:
@@ -123,6 +123,7 @@ const ADMIN_PROVIDER_QUERY_NO_ACTIVE_TEST_CANDIDATE_DETAIL: &str =
|
||||
"No active endpoint or API key found";
|
||||
const ADMIN_PROVIDER_QUERY_INVALID_MAPPED_MODEL_DETAIL: &str =
|
||||
"mapped_model_name is not valid for the selected model and endpoint";
|
||||
const PROVIDER_QUERY_KEY_MODEL_NOT_ALLOWED_SKIP_REASON: &str = "key_model_not_allowed";
|
||||
const ANTIGRAVITY_PROVIDER_CACHE_KEY_PREFIX: &str = "upstream_models_provider:";
|
||||
const DEFAULT_PROVIDER_QUERY_TEST_MESSAGE: &str = "Hello! This is a test message.";
|
||||
static PROVIDER_QUERY_POOL_LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
|
||||
@@ -933,6 +934,35 @@ fn provider_query_selected_key_ids_all_exist(
|
||||
.all(|id| keys.iter().any(|key| key.id == *id))
|
||||
}
|
||||
|
||||
fn provider_query_model_name_matches(left: &str, right: &str) -> bool {
|
||||
let left = left.trim();
|
||||
let right = right.trim();
|
||||
!left.is_empty() && !right.is_empty() && left.eq_ignore_ascii_case(right)
|
||||
}
|
||||
|
||||
fn provider_query_key_allows_effective_test_model(
|
||||
key: &StoredProviderCatalogKey,
|
||||
requested_model: &str,
|
||||
effective_model: &str,
|
||||
) -> bool {
|
||||
let allowed_models = json_string_list(key.allowed_models.as_ref());
|
||||
if key.allowed_models.is_none() || allowed_models.is_empty() {
|
||||
return true;
|
||||
}
|
||||
|
||||
let requested_base_model = crate::ai_serving::model_directive_base_model(requested_model);
|
||||
allowed_models
|
||||
.iter()
|
||||
.map(String::as_str)
|
||||
.any(|allowed_model| {
|
||||
provider_query_model_name_matches(allowed_model, requested_model)
|
||||
|| provider_query_model_name_matches(allowed_model, effective_model)
|
||||
|| requested_base_model.as_deref().is_some_and(|base_model| {
|
||||
provider_query_model_name_matches(allowed_model, base_model)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn provider_query_test_key_sort_key(
|
||||
provider_type: &str,
|
||||
key: &StoredProviderCatalogKey,
|
||||
@@ -1382,17 +1412,43 @@ async fn provider_query_build_kiro_test_candidates(
|
||||
.unwrap_or(requested_model.clone())
|
||||
};
|
||||
|
||||
let mut keys = all_keys
|
||||
let now_unix_secs = current_unix_ms() / 1000;
|
||||
let mut keys = Vec::new();
|
||||
let mut model_skipped_candidates = Vec::new();
|
||||
|
||||
for key in all_keys
|
||||
.into_iter()
|
||||
.filter(|key| key.is_active)
|
||||
.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)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let now_unix_secs = current_unix_ms() / 1000;
|
||||
{
|
||||
if provider_query_key_allows_effective_test_model(&key, &requested_model, &effective_model)
|
||||
{
|
||||
keys.push(key);
|
||||
} else {
|
||||
model_skipped_candidates.push(ProviderQueryTestCandidate {
|
||||
endpoint: endpoint.clone(),
|
||||
key,
|
||||
effective_model: effective_model.clone(),
|
||||
scheduler_skip_reason: Some(
|
||||
PROVIDER_QUERY_KEY_MODEL_NOT_ALLOWED_SKIP_REASON.to_string(),
|
||||
),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let candidates = if test_mode.eq_ignore_ascii_case("pool") {
|
||||
model_skipped_candidates.sort_by_key(|candidate| {
|
||||
provider_query_test_key_sort_key(
|
||||
provider.provider_type.as_str(),
|
||||
&candidate.key,
|
||||
&endpoint.api_format,
|
||||
now_unix_secs,
|
||||
)
|
||||
});
|
||||
|
||||
let scheduled_candidates = if test_mode.eq_ignore_ascii_case("pool") {
|
||||
if let Some(pool_config) =
|
||||
admin_provider_pool_config_from_config_value(provider.config.as_ref())
|
||||
{
|
||||
@@ -1442,6 +1498,8 @@ async fn provider_query_build_kiro_test_candidates(
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
};
|
||||
let mut candidates = model_skipped_candidates;
|
||||
candidates.extend(scheduled_candidates);
|
||||
|
||||
if candidates.is_empty() {
|
||||
return Err(build_admin_provider_query_not_found_response(
|
||||
|
||||
@@ -63,6 +63,68 @@ fn sample_openai_image_transport(provider_type: &str) -> AdminGatewayProviderTra
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_catalog_key_with_allowed_models(
|
||||
allowed_models: Option<serde_json::Value>,
|
||||
) -> aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey {
|
||||
let mut key =
|
||||
aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey::new(
|
||||
"key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"key".to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("sample provider key should build");
|
||||
key.allowed_models = allowed_models;
|
||||
key
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_query_model_test_allows_keys_without_model_restrictions() {
|
||||
let unrestricted = sample_catalog_key_with_allowed_models(None);
|
||||
let empty = sample_catalog_key_with_allowed_models(Some(json!([])));
|
||||
|
||||
assert!(provider_query_key_allows_effective_test_model(
|
||||
&unrestricted,
|
||||
"model-b",
|
||||
"model-b-upstream",
|
||||
));
|
||||
assert!(provider_query_key_allows_effective_test_model(
|
||||
&empty,
|
||||
"model-b",
|
||||
"model-b-upstream",
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_query_model_test_filters_key_disallowed_for_requested_model() {
|
||||
let key = sample_catalog_key_with_allowed_models(Some(json!(["model-a"])));
|
||||
|
||||
assert!(!provider_query_key_allows_effective_test_model(
|
||||
&key,
|
||||
"model-b",
|
||||
"model-b-upstream",
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_query_model_test_allows_key_for_requested_or_mapped_model() {
|
||||
let requested_allowed = sample_catalog_key_with_allowed_models(Some(json!(["model-b"])));
|
||||
let mapped_allowed = sample_catalog_key_with_allowed_models(Some(json!(["MODEL-B-UPSTREAM"])));
|
||||
|
||||
assert!(provider_query_key_allows_effective_test_model(
|
||||
&requested_allowed,
|
||||
"model-b",
|
||||
"model-b-upstream",
|
||||
));
|
||||
assert!(provider_query_key_allows_effective_test_model(
|
||||
&mapped_allowed,
|
||||
"model-b",
|
||||
"model-b-upstream",
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_query_test_request_body_preserves_custom_model() {
|
||||
let payload = json!({
|
||||
|
||||
+5
@@ -372,6 +372,11 @@ describe('isModelTestableEndpoint', () => {
|
||||
})
|
||||
|
||||
describe('formatModelTestDiagnostic', () => {
|
||||
it('maps model permission skips to an actionable label', () => {
|
||||
expect(formatModelTestDiagnostic('key_model_not_allowed'))
|
||||
.toBe('Key 未允许当前模型,已跳过')
|
||||
})
|
||||
|
||||
it('maps pool account blocked scheduler code to an actionable label', () => {
|
||||
expect(formatModelTestDiagnostic('pool_account_blocked')).toBe('账号已失效,需重新授权')
|
||||
})
|
||||
|
||||
@@ -41,6 +41,7 @@ const MODEL_TEST_BEARER_INHERITS_PROVIDER_FORMATS = new Set([
|
||||
])
|
||||
|
||||
const MODEL_TEST_DIAGNOSTIC_LABELS: Record<string, string> = {
|
||||
key_model_not_allowed: 'Key 未允许当前模型,已跳过',
|
||||
pool_account_blocked: '账号已失效,需重新授权',
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user