This commit is contained in:
fawney19
2026-05-19 03:16:35 +08:00
77 changed files with 4051 additions and 333 deletions

View File

@@ -18,6 +18,7 @@ use crate::ai_serving::{
};
use crate::clock::current_unix_ms;
use crate::execution_runtime;
use crate::handlers::admin::provider::write::provider::reconcile_admin_fixed_provider_template_endpoints;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::handlers::shared::provider_pool::{
admin_provider_pool_config_from_config_value, read_admin_provider_pool_runtime_state,
@@ -685,7 +686,6 @@ fn provider_query_build_test_request_body_with_model_policy(
"content": provider_query_extract_message(payload)
.unwrap_or_else(|| DEFAULT_PROVIDER_QUERY_TEST_MESSAGE.to_string())
}],
"max_tokens": 30,
"temperature": 0.7,
"stream": true,
})
@@ -1219,6 +1219,8 @@ async fn provider_query_build_kiro_test_candidates(
payload: &Value,
requested_model_override: Option<&str>,
) -> Result<Vec<ProviderQueryTestCandidate>, Response<Body>> {
provider_query_reconcile_fixed_provider_endpoints_for_test_model(state, provider).await?;
let provider_ids = vec![provider.id.clone()];
let endpoints = state
.app()
@@ -1424,6 +1426,33 @@ async fn provider_query_build_kiro_test_candidates(
Ok(candidates)
}
async fn provider_query_reconcile_fixed_provider_endpoints_for_test_model(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
) -> Result<(), Response<Body>> {
if state
.fixed_provider_template(&provider.provider_type)
.is_none()
|| !state.has_provider_catalog_data_writer()
{
return Ok(());
}
reconcile_admin_fixed_provider_template_endpoints(state, provider)
.await
.map_err(|err| {
warn!(
provider_id = %provider.id,
provider_type = %provider.provider_type,
error = ?err,
"admin provider-query test-model: failed to reconcile fixed provider endpoints"
);
build_admin_provider_query_bad_request_response(
ADMIN_PROVIDER_QUERY_NO_ACTIVE_API_KEY_DETAIL,
)
})
}
fn provider_query_decode_execution_body(
result: &aether_contracts::ExecutionResult,
) -> Option<Vec<u8>> {
@@ -1453,7 +1482,7 @@ fn provider_query_standard_execution_response_body(
provider_api_format: &str,
result: &aether_contracts::ExecutionResult,
) -> Option<Value> {
result
let body = result
.body
.as_ref()
.and_then(|body| body.json_body.clone())
@@ -1461,7 +1490,15 @@ fn provider_query_standard_execution_response_body(
provider_query_decode_execution_body(result).and_then(|body| {
provider_query_aggregate_standard_stream_sync_response(provider_api_format, &body)
})
})
})?;
if result.status_code < 400
&& provider_query_normalize_api_format_alias(provider_api_format)
== "gemini:generate_content"
&& aether_ai_formats::formats::gemini::generate_content::response::from_raw(&body).is_none()
{
return None;
}
Some(body)
}
fn provider_query_extract_error_message(
@@ -2705,7 +2742,7 @@ async fn provider_query_execute_standard_test_candidate(
}
"openai:embedding" | "gemini:embedding" | "jina:embedding" | "doubao:embedding"
| "openai:rerank" | "jina:rerank" => {
let Some(provider_request_body) =
let Some(mut provider_request_body) =
crate::ai_serving::build_standard_request_body_with_model_directives_and_request_headers(
&request_body,
client_api_format,
@@ -2725,6 +2762,18 @@ async fn provider_query_execute_standard_test_candidate(
format!("Provider request body could not be built for {provider_api_format}"),
));
};
if let Err(err) = crate::provider_transport::apply_transport_request_body_semantics(
&mut provider_request_body,
&transport,
normalized_provider_api_format.as_str(),
) {
return Ok(provider_query_skipped_execution_outcome(
provider_request_body,
format!(
"Provider request body is not compatible with transport semantics: {err}"
),
));
}
provider_request_body
}
_ => {
@@ -2805,7 +2854,7 @@ async fn provider_query_execute_standard_test_candidate(
*synthetic_request.headers_mut() = incoming_request_headers;
let (parts, _) = synthetic_request.into_parts();
let request_url = crate::provider_transport::build_transport_request_url(
let request_url = crate::provider_transport::build_transport_request_url_for_request_body(
&transport,
crate::provider_transport::TransportRequestUrlParams {
provider_api_format,
@@ -2814,6 +2863,7 @@ async fn provider_query_execute_standard_test_candidate(
request_query: parts.uri.query(),
kiro_api_region: None,
},
Some(&provider_request_body),
);
let Some(request_url) = request_url else {
return Ok(provider_query_skipped_execution_outcome(

View File

@@ -63,10 +63,10 @@ pub(super) fn provider_query_standard_test_unsupported_reason(
api_format,
)
}
"gemini:generate_content"
if crate::provider_transport::is_vertex_api_key_transport_context(transport) =>
"gemini:generate_content" | "gemini:embedding"
if crate::provider_transport::is_vertex_transport_context(transport) =>
{
aether_provider_transport::vertex::local_vertex_api_key_gemini_transport_unsupported_reason_with_network(
aether_provider_transport::vertex::local_vertex_gemini_transport_unsupported_reason_with_network(
transport,
)
}

View File

@@ -9,11 +9,29 @@ pub(super) fn provider_query_test_attempt_payload(
candidate: &ProviderQueryTestCandidate,
execution: &ProviderQueryExecutionOutcome,
) -> Value {
let endpoint_route = provider_query_endpoint_route_payload(candidate, execution);
let endpoint_product = endpoint_route
.get("product")
.cloned()
.unwrap_or(Value::Null);
let endpoint_variant = endpoint_route
.get("variant")
.cloned()
.unwrap_or(Value::Null);
let endpoint_action = endpoint_route.get("action").cloned().unwrap_or(Value::Null);
let endpoint_batch_strategy = endpoint_route
.get("batch_strategy")
.cloned()
.unwrap_or(Value::Null);
json!({
"candidate_index": candidate_index,
"retry_index": 0,
"endpoint_api_format": candidate.endpoint.api_format,
"endpoint_base_url": candidate.endpoint.base_url,
"endpoint_product": endpoint_product,
"endpoint_variant": endpoint_variant,
"endpoint_action": endpoint_action,
"endpoint_batch_strategy": endpoint_batch_strategy,
"key_name": provider_query_key_display_name(&candidate.key),
"key_id": candidate.key.id,
"auth_type": candidate.key.auth_type,
@@ -31,6 +49,129 @@ pub(super) fn provider_query_test_attempt_payload(
})
}
fn provider_query_endpoint_route_payload(
candidate: &ProviderQueryTestCandidate,
execution: &ProviderQueryExecutionOutcome,
) -> Value {
let api_format = aether_ai_formats::normalize_api_format_alias(&candidate.endpoint.api_format);
let request_url = execution.request_url.to_ascii_lowercase();
let base_url = candidate.endpoint.base_url.to_ascii_lowercase();
let is_vertex = request_url.contains("aiplatform.googleapis.com")
|| base_url.contains("aiplatform.googleapis.com");
let is_gemini_api = request_url.contains("generativelanguage.googleapis.com")
|| base_url.contains("generativelanguage.googleapis.com");
let is_openai_compat =
request_url.contains("/endpoints/openapi") || request_url.contains("/openai/");
let is_batch = execution
.request_body
.get("requests")
.and_then(Value::as_array)
.is_some_and(|items| !items.is_empty());
let vertex_instance_count = execution
.request_body
.get("instances")
.and_then(Value::as_array)
.map(Vec::len)
.unwrap_or(0);
let (product, variant, action, batch_strategy) = match api_format.as_str() {
"gemini:embedding" if is_vertex => (
"Vertex AI",
"vertex_native",
"predict",
if vertex_instance_count > 1 {
"predict_instances"
} else {
"single_instance"
},
),
"gemini:embedding" if is_gemini_api => (
"Gemini API",
"gemini_native",
if is_batch {
"batchEmbedContents"
} else {
"embedContent"
},
if is_batch {
"native_batch"
} else {
"single_native"
},
),
"gemini:embedding" => (
"Gemini native",
"gemini_native",
if is_batch {
"batchEmbedContents"
} else {
"embedContent"
},
if is_batch {
"native_batch"
} else {
"single_native"
},
),
"gemini:generate_content" if is_vertex => {
("Vertex AI", "vertex_native", "generateContent", "")
}
"gemini:generate_content" if is_gemini_api => {
("Gemini API", "gemini_native", "generateContent", "")
}
"gemini:generate_content" => ("Gemini native", "gemini_native", "generateContent", ""),
"openai:embedding" if is_vertex && is_openai_compat => (
"Vertex AI OpenAI-compatible",
"openai_compatible",
"embeddings",
"openai_batch",
),
"openai:embedding" if is_gemini_api && is_openai_compat => (
"Gemini API OpenAI-compatible",
"openai_compatible",
"embeddings",
"openai_batch",
),
"openai:embedding" => (
"OpenAI-compatible",
"openai_compatible",
"embeddings",
"openai_batch",
),
"openai:chat" if is_vertex && is_openai_compat => (
"Vertex AI OpenAI-compatible",
"openai_compatible",
"chat/completions",
"",
),
"openai:chat" if is_gemini_api && is_openai_compat => (
"Gemini API OpenAI-compatible",
"openai_compatible",
"chat/completions",
"",
),
"openai:chat" => (
"OpenAI-compatible",
"openai_compatible",
"chat/completions",
"",
),
_ => (
"Provider endpoint",
"provider_native",
"provider_request",
"",
),
};
json!({
"product": product,
"variant": variant,
"action": action,
"batch_strategy": batch_strategy,
})
}
fn provider_query_redact_diagnostic_headers(
headers: &BTreeMap<String, String>,
) -> BTreeMap<String, String> {

View File

@@ -125,6 +125,17 @@ fn provider_query_test_request_body_keeps_non_empty_conversation() {
);
}
#[test]
fn provider_query_default_test_request_body_does_not_set_max_tokens() {
let body = provider_query_build_test_request_body(&json!({}), "fallback-model");
assert_eq!(body["model"], json!("fallback-model"));
assert!(
body.get("max_tokens").is_none(),
"admin model test must not silently force a low max_tokens value"
);
}
#[test]
fn provider_query_failover_request_body_overrides_custom_model() {
let payload = json!({
@@ -313,6 +324,40 @@ fn provider_query_responses_test_request_body_defaults_to_responses_input() {
assert!(body.get("messages").is_none());
}
#[test]
fn provider_query_standard_test_rejects_gemini_success_without_visible_output() {
let result = aether_contracts::ExecutionResult {
request_id: "provider-test".to_string(),
candidate_id: Some("candidate-0".to_string()),
status_code: 200,
headers: BTreeMap::new(),
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"candidates": [{
"content": {"role": "model"},
"finishReason": "MAX_TOKENS"
}],
"usageMetadata": {
"promptTokenCount": 8,
"candidatesTokenCount": 1,
"thoughtsTokenCount": 25,
"totalTokenCount": 34
},
"modelVersion": "gemini-3-flash-preview",
"responseId": "resp-empty"
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
};
assert!(
provider_query_standard_execution_response_body("gemini:generate_content", &result)
.is_none()
);
}
#[test]
fn provider_query_test_adapter_routes_fixed_provider_endpoint_types() {
assert_eq!(

View File

@@ -161,8 +161,12 @@ pub(crate) fn validate_vertex_api_formats(
}
let allowed = match auth_type {
"api_key" => &["gemini:generate_content"][..],
"service_account" | "vertex_ai" => &["claude:messages", "gemini:generate_content"][..],
"api_key" => &["gemini:generate_content", "gemini:embedding"][..],
"service_account" | "vertex_ai" => &[
"claude:messages",
"gemini:generate_content",
"gemini:embedding",
][..],
_ => return Ok(()),
};
let invalid = api_formats
@@ -375,4 +379,27 @@ mod tests {
)
.is_err());
}
#[test]
fn validate_vertex_api_formats_allows_gemini_embedding() {
assert!(validate_vertex_api_formats(
"vertex_ai",
"api_key",
&[
"gemini:generate_content".to_string(),
"gemini:embedding".to_string()
],
)
.is_ok());
assert!(validate_vertex_api_formats(
"vertex_ai",
"service_account",
&[
"claude:messages".to_string(),
"gemini:generate_content".to_string(),
"gemini:embedding".to_string()
],
)
.is_ok());
}
}

View File

@@ -169,7 +169,11 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
}
let mut provider_request_body = match format_value.as_str() {
"openai:chat" | "claude:messages" => json!({
"openai:chat" => json!({
"model": model,
"messages": [{"role": "user", "content": "Health check"}],
}),
"claude:messages" => json!({
"model": model,
"messages": [{"role": "user", "content": "Health check"}],
"max_tokens": 5,
@@ -179,9 +183,6 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
"role": "user",
"parts": [{"text": "Health check"}],
}],
"generationConfig": {
"maxOutputTokens": 5,
},
}),
_ => return None,
};