i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-01 17:00:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(gateway): normalize Gemini Vertex embedding transport

Browse Source
This commit is contained in:
MMEXA
2026-05-18 15:39:29 +00:00
parent 66f154a251
commit b480f3aaff
24 changed files with 1046 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs
View File

@@ -241,11 +241,29 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
apply_transport_request_body_semantics(
if let Err(err) = apply_transport_request_body_semantics(
&mut provider_request_body,
transport,
provider_api_format,
);
) {
mark_skipped_local_standard_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_request_body_semantics_failed",
CandidateFailureDiagnostic::request_conversion_failed(
spec_metadata.api_format,
provider_api_format,
"standard_family_transport_body_semantics",
err.to_string(),
),
)
.await;
return None;
}
if let Some(mapping) =
crate::system_features::reasoning_model_directive_mapping_for_api_format_and_model(
state,
@@ -266,11 +284,29 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
upstream_is_stream,
request_requires_body_stream_field(body_json, force_body_stream_field),
);
apply_transport_request_body_semantics(
if let Err(err) = apply_transport_request_body_semantics(
&mut provider_request_body,
transport,
provider_api_format,
);
) {
mark_skipped_local_standard_candidate_with_failure_diagnostic(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_request_body_semantics_failed",
CandidateFailureDiagnostic::request_conversion_failed(
spec_metadata.api_format,
provider_api_format,
"standard_family_transport_body_semantics_after_model_directives",
err.to_string(),
),
)
.await;
return None;
}
}
if let Some(kiro_auth) = kiro_auth.as_ref() {
@@ -382,22 +418,12 @@ fn apply_transport_request_body_semantics(
provider_request_body: &mut Value,
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
) {
if !crate::ai_serving::api_format_alias_matches(provider_api_format, "gemini:embedding")
|| !crate::ai_serving::transport::vertex::is_vertex_transport_context(transport)
{
return;
}
let Some(object) = provider_request_body.as_object_mut() else {
return;
};
if object.contains_key("requests") {
return;
}
object.remove("model");
) -> Result<(), crate::ai_serving::transport::TransportRequestBodySemanticsError> {
crate::ai_serving::transport::apply_transport_request_body_semantics(
provider_request_body,
transport,
provider_api_format,
)
}
async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(

6
apps/aether-gateway/src/ai_serving/transport.rs
View File

@@ -51,7 +51,8 @@ pub(crate) use aether_provider_transport::{
apply_local_body_rules_with_request_headers, apply_local_header_rules,
apply_local_header_rules_with_request_headers, apply_standard_provider_request_body_rules,
apply_standard_provider_request_body_rules_with_request_headers,
body_rules_are_locally_supported, body_rules_handle_path, body_rules_have_enabled_rules,
apply_transport_request_body_semantics, body_rules_are_locally_supported,
body_rules_handle_path, body_rules_have_enabled_rules,
build_cross_format_openai_chat_upstream_url, build_cross_format_openai_responses_upstream_url,
build_gemini_files_headers, build_gemini_files_request_body, build_gemini_files_upstream_url,
build_kiro_cross_format_upstream_url, build_local_openai_chat_upstream_url,
@@ -89,5 +90,6 @@ pub(crate) use aether_provider_transport::{
SameFormatProviderRequestBehaviorParams, SameFormatProviderRequestBodyInput,
SameFormatProviderUpstreamUrlParams, StandardPlanFallbackAcceptPolicy,
StandardPlanFallbackHeadersInput, StandardProviderRequestHeaders,
StandardProviderRequestHeadersInput, TransportRequestUrlParams,
StandardProviderRequestHeadersInput, TransportRequestBodySemanticsError,
TransportRequestUrlParams,
};

2
apps/aether-gateway/src/api/ai/gemini.rs
View File

@@ -11,7 +11,7 @@ pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> {
pub(crate) fn local_path(api_format: &str) -> Option<&'static str> {
match crate::ai_serving::normalize_api_format_alias(api_format).as_str() {
"gemini" | "gemini:generate_content" => Some("/v1beta/models/{model}:{action}"),
"gemini:embedding" => Some("/v1/embeddings"),
"gemini:embedding" => Some("/v1beta/models/{model}:{action}"),
"gemini:video" => Some("/v1beta/models/{model}:predictLongRunning"),
"gemini:files" => Some("/v1beta/files"),
_ => None,

7
apps/aether-gateway/src/api/ai/registry.rs
View File

@@ -86,7 +86,12 @@ mod tests {
fn supports_data_api_endpoint_signatures_and_public_paths() {
for (api_format, family, kind, path) in [
("openai:embedding", "openai", "embedding", "/v1/embeddings"),
("gemini:embedding", "gemini", "embedding", "/v1/embeddings"),
(
"gemini:embedding",
"gemini",
"embedding",
"/v1beta/models/{model}:{action}",
),
("jina:embedding", "jina", "embedding", "/v1/embeddings"),
("doubao:embedding", "doubao", "embedding", "/v1/embeddings"),
("openai:rerank", "openai", "rerank", "/v1/rerank"),

79
apps/aether-gateway/src/execution_runtime/fallback.rs
View File

@@ -2,7 +2,8 @@ use aether_contracts::{ExecutionPlan, ExecutionResult};
use serde_json::Value;
use crate::orchestration::{
resolve_local_failover_analysis_for_attempt, LocalFailoverAnalysis, LocalFailoverDecision,
resolve_local_failover_analysis_for_attempt, LocalFailoverAnalysis,
LocalFailoverClassification, LocalFailoverDecision,
};
use crate::AppState;
@@ -48,6 +49,15 @@ pub(crate) async fn analyze_local_candidate_failover_sync(
return LocalFailoverAnalysis::use_default();
}
if let Some(error) = result.error.as_ref() {
if !error.retryable && !error.failover_recommended {
return LocalFailoverAnalysis {
classification: LocalFailoverClassification::StopExecutionError,
decision: LocalFailoverDecision::StopLocalFailover,
};
}
}
resolve_local_failover_analysis_for_attempt(
state,
plan,
@@ -326,14 +336,14 @@ pub(crate) fn resolve_core_stream_direct_finalize_report_kind(plan_kind: &str) -
mod tests {
use std::collections::BTreeSet;
use aether_contracts::ExecutionResult;
use aether_contracts::{ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionResult};
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use super::{
resolve_core_stream_error_finalize_report_kind,
analyze_local_candidate_failover_sync, resolve_core_stream_error_finalize_report_kind,
resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_stream,
should_fallback_to_control_sync, should_retry_next_local_candidate_stream,
should_retry_next_local_candidate_sync, should_stop_local_candidate_failover_stream,
@@ -607,6 +617,69 @@ mod tests {
);
}
#[tokio::test]
async fn sync_failover_honors_non_retryable_execution_error() {
let result = ExecutionResult {
request_id: "req-1".to_string(),
candidate_id: None,
status_code: 502,
headers: Default::default(),
body: None,
telemetry: None,
error: Some(ExecutionError {
kind: ExecutionErrorKind::Upstream5xx,
phase: ExecutionPhase::Finalize,
message: "provider returned HTTP 200 without visible model output".to_string(),
upstream_status: Some(200),
retryable: false,
failover_recommended: false,
}),
};
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
let state = build_state_with_provider_config(None);
let plan = sample_plan();
let analysis = analyze_local_candidate_failover_sync(
&state,
&plan,
"openai_chat_sync",
Some(&local_report_context),
&result,
Some("provider returned HTTP 200 without visible model output"),
)
.await;
assert_eq!(
analysis.decision,
crate::orchestration::LocalFailoverDecision::StopLocalFailover
);
assert!(
!should_retry_next_local_candidate_sync(
&state,
&plan,
"openai_chat_sync",
Some(&local_report_context),
&result,
Some("provider returned HTTP 200 without visible model output"),
)
.await
);
assert!(
should_stop_local_candidate_failover_sync(
&state,
&plan,
"openai_chat_sync",
Some(&local_report_context),
&result,
Some("provider returned HTTP 200 without visible model output"),
)
.await
);
}
#[tokio::test]
async fn sync_retry_next_candidate_skips_video_follow_up_plan_kinds() {
let result = ExecutionResult {

17
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs
View File

@@ -2304,7 +2304,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,
@@ -2324,6 +2324,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
}
_ => {
@@ -2404,7 +2416,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,
@@ -2413,6 +2425,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(

6
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/adapter.rs
View File

@@ -62,10 +62,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,
)
}

141
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/summary.rs
View File

@@ -7,11 +7,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,
@@ -29,6 +47,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,
})
}
pub(super) fn provider_query_candidate_summary_payload(
total_candidates: usize,
total_attempts: usize,

2
apps/aether-gateway/src/orchestration/classifier.rs
View File

@@ -32,6 +32,7 @@ pub(crate) enum LocalFailoverClassification {
UseDefault,
StopStatusCode,
StopErrorPattern,
StopExecutionError,
RetrySuccessPattern,
RetryStatusCode,
RetryUpstreamFailure,
@@ -43,6 +44,7 @@ impl LocalFailoverClassification {
Self::UseDefault => "use_default",
Self::StopStatusCode => "stop_status_code",
Self::StopErrorPattern => "stop_error_pattern",
Self::StopExecutionError => "stop_execution_error",
Self::RetrySuccessPattern => "retry_success_pattern",
Self::RetryStatusCode => "retry_status_code",
Self::RetryUpstreamFailure => "retry_upstream_failure",

3
apps/aether-gateway/src/orchestration/effects.rs
View File

@@ -821,7 +821,8 @@ fn local_candidate_failure_should_invalidate_affinity(
LocalFailoverClassification::UseDefault | LocalFailoverClassification::StopStatusCode => {
status_code >= 500
}
LocalFailoverClassification::StopErrorPattern => false,
LocalFailoverClassification::StopErrorPattern
| LocalFailoverClassification::StopExecutionError => false,
}
}

3
apps/aether-gateway/src/orchestration/health.rs
View File

@@ -155,7 +155,8 @@ fn local_candidate_failure_should_project_health(
LocalFailoverClassification::UseDefault | LocalFailoverClassification::StopStatusCode => {
status_code >= 500
}
LocalFailoverClassification::StopErrorPattern => false,
LocalFailoverClassification::StopErrorPattern
| LocalFailoverClassification::StopExecutionError => false,
}
}

4
apps/aether-gateway/src/orchestration/mod.rs
View File

@@ -92,7 +92,9 @@ pub(crate) fn build_local_error_flow_metadata(
) -> Value {
let safe_to_expose = matches!(
analysis.classification,
LocalFailoverClassification::StopStatusCode | LocalFailoverClassification::StopErrorPattern
LocalFailoverClassification::StopStatusCode
| LocalFailoverClassification::StopErrorPattern
| LocalFailoverClassification::StopExecutionError
);
let propagation = match analysis.decision {
LocalFailoverDecision::RetryNextCandidate => "suppressed",

5
apps/aether-gateway/src/orchestration/recovery.rs
View File

@@ -57,7 +57,10 @@ const fn decision_from_classification(
match classification {
LocalFailoverClassification::UseDefault => LocalFailoverDecision::UseDefault,
LocalFailoverClassification::StopStatusCode
| LocalFailoverClassification::StopErrorPattern => LocalFailoverDecision::StopLocalFailover,
| LocalFailoverClassification::StopErrorPattern
| LocalFailoverClassification::StopExecutionError => {
LocalFailoverDecision::StopLocalFailover
}
LocalFailoverClassification::RetrySuccessPattern
| LocalFailoverClassification::RetryStatusCode
| LocalFailoverClassification::RetryUpstreamFailure => {

140
apps/aether-gateway/src/tests/control/admin/provider_query.rs
View File

@@ -1256,6 +1256,146 @@ async fn gateway_handles_admin_provider_query_gemini_embedding_model_test() {
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_vertex_gemini_embedding_model_test() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-vertex-ai");
assert_eq!(plan.endpoint_id, "endpoint-vertex-gemini-embedding");
assert_eq!(plan.key_id, "key-vertex-gemini-embedding");
assert_eq!(plan.client_api_format, "openai:embedding");
assert_eq!(plan.provider_api_format, "gemini:embedding");
assert_eq!(
plan.url,
"https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-embedding-2:predict?key=sk-vertex-gemini-embedding"
);
assert_eq!(plan.model_name.as_deref(), Some("gemini-embedding-2"));
assert!(!plan.stream);
let body = plan.body.json_body.as_ref().expect("json body");
assert!(
body.get("model").is_none(),
"Vertex predict carries the model in the URL path; the test body must not repeat it"
);
assert_eq!(
body["instances"][0]["content"],
json!("This is a test embedding input.")
);
assert!(body.get("content").is_none());
assert!(body.get("requests").is_none());
assert!(
body.get("stream").is_none(),
"gemini embedding provider body must not carry stream"
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"predictions": [
{
"embeddings": {
"values": [0.1, 0.2, 0.3]
}
}
],
"deployedModelId": "gemini-embedding-2"
}
},
"telemetry": {
"elapsed_ms": 27
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-vertex-ai", "Vertex AI", 10);
provider.provider_type = "vertex_ai".to_string();
let mut key = sample_key(
"key-vertex-gemini-embedding",
"provider-vertex-ai",
"gemini:embedding",
"sk-vertex-gemini-embedding",
);
key.allowed_models = Some(json!(["gemini-embedding-2"]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-vertex-gemini-embedding",
"provider-vertex-ai",
"gemini:embedding",
"https://aiplatform.googleapis.com",
)],
vec![key],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/provider-query/test-model"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-vertex-ai",
"model": "gemini-embedding-2",
"api_format": "gemini:embedding",
"endpoint_id": "endpoint-vertex-gemini-embedding",
"request_body": {
"model": "gemini-embedding-2",
"input": "This is a test embedding input."
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["error"], serde_json::Value::Null);
assert_eq!(payload["attempts"][0]["status"], json!("success"));
assert_eq!(
payload["attempts"][0]["request_body"]["instances"][0]["content"],
json!("This is a test embedding input.")
);
assert_eq!(
payload["attempts"][0]["endpoint_product"],
json!("Vertex AI")
);
assert_eq!(
payload["attempts"][0]["endpoint_variant"],
json!("vertex_native")
);
assert_eq!(payload["attempts"][0]["endpoint_action"], json!("predict"));
assert_eq!(
payload["attempts"][0]["endpoint_batch_strategy"],
json!("single_instance")
);
assert!(
payload["attempts"][0]["request_body"]
.get("model")
.is_none(),
"attempt debug payload must expose the exact Vertex body without a duplicate model"
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_jina_embedding_model_test() {
let execution_runtime = Router::new().route(

8
apps/aether-gateway/src/tests/control/admin/system.rs
View File

@@ -1152,6 +1152,14 @@ async fn gateway_handles_admin_system_api_formats_locally_with_trusted_admin_pri
.expect("formats should be an array");
assert_eq!(formats[0]["value"], "openai:chat");
assert_eq!(formats[0]["default_path"], "/v1/chat/completions");
let gemini_embedding = formats
.iter()
.find(|item| item["value"] == "gemini:embedding")
.expect("gemini embedding format should exist");
assert_eq!(
gemini_embedding["default_path"],
"/v1beta/models/{model}:{action}"
);
assert!(formats
.iter()
.any(|item| item["value"] == "openai:embedding"));

35
apps/aether-gateway/src/tests/control/proxy/embeddings.rs
View File

@@ -194,7 +194,7 @@ fn vertex_gemini_embedding_conversion_execution_runtime() -> Router {
"/v1/execute/sync",
any(|Json(plan): Json<ExecutionPlan>| async move {
assert_openai_to_vertex_gemini_embedding_execution_plan(&plan);
Json(gemini_embedding_execution_result(&plan))
Json(vertex_gemini_embedding_execution_result(&plan))
}),
)
}
@@ -342,7 +342,7 @@ fn assert_openai_to_vertex_gemini_embedding_execution_plan(plan: &ExecutionPlan)
assert_eq!(plan.method, "POST");
assert_eq!(
plan.url,
"https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-embedding-2:embedContent?key=sk-upstream-vertex-gemini-embedding"
"https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-embedding-2:predict?key=sk-upstream-vertex-gemini-embedding"
);
assert_eq!(
plan.model_name.as_deref(),
@@ -352,9 +352,10 @@ fn assert_openai_to_vertex_gemini_embedding_execution_plan(plan: &ExecutionPlan)
let body = plan.body.json_body.as_ref().expect("json request body");
assert!(
body.get("model").is_none(),
"Vertex embedContent carries the model in the path; the body must not repeat it"
"Vertex predict carries the model in the path; the body must not repeat it"
);
assert_eq!(body["content"]["parts"][0]["text"], "hello");
assert_eq!(body["instances"][0]["content"], "hello");
assert!(body.get("content").is_none());
assert!(body.get("input").is_none());
assert!(body.get("messages").is_none());
}
@@ -453,6 +454,30 @@ fn gemini_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
}
}
fn vertex_gemini_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
ExecutionResult {
request_id: plan.request_id.clone(),
candidate_id: plan.candidate_id.clone(),
status_code: 200,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
body: Some(ResponseBody {
json_body: Some(json!({
"predictions": [
{
"embeddings": {
"values": [0.1, 0.2, 0.3]
}
}
],
"deployedModelId": "gemini-embedding-2"
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
}
}
fn gemini_batch_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
ExecutionResult {
request_id: plan.request_id.clone(),
@@ -636,7 +661,7 @@ async fn embeddings_route_converts_openai_payload_to_vertex_gemini_embedding_pro
assert_eq!(endpoint_signature.as_deref(), Some("openai:embedding"));
let payload: serde_json::Value = serde_json::from_str(&body_text).expect("body should parse");
assert_eq!(payload["object"], "list");
assert_eq!(payload["model"], "gemini-embedding-2-preview");
assert_eq!(payload["model"], "gemini-embedding-2");
assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3]));
gateway_handle.abort();