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

fix(gateway): route Gemini embedding batches correctly

Browse Source
This commit is contained in:
MMEXA
2026-05-17 14:24:32 +00:00
parent a2f91b4108
commit 81ff375bfd
37 changed files with 1353 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
apps/aether-gateway/src/ai_serving/api.rs
View File

@@ -49,14 +49,15 @@ pub(crate) use aether_ai_formats::api::{
resolve_claude_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
resolve_local_image_stream_spec, resolve_local_image_sync_spec,
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
AiControlPlanRequest, ExecutionRuntimeAuthContext, LocalCoreSyncErrorKind,
LocalOpenAiImageSpec, LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec,
StreamingStandardTerminalObserver, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_FILES_DOWNLOAD_PLAN_KIND,
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
resolve_openai_embedding_sync_spec, AiControlPlanRequest, ExecutionRuntimeAuthContext,
LocalCoreSyncErrorKind, LocalOpenAiImageSpec, LocalSameFormatProviderFamily,
LocalSameFormatProviderSpec, LocalStandardSourceFamily, LocalStandardSourceMode,
LocalStandardSpec, StreamingStandardTerminalObserver, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_EMBEDDING_SYNC_PLAN_KIND,
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
};

25
apps/aether-gateway/src/ai_serving/mod.rs
View File

@@ -64,7 +64,8 @@ pub(crate) use self::transport::{
candidate_common_transport_skip_reason, candidate_transport_pair_skip_reason,
request_conversion_direct_auth, request_conversion_enabled_for_transport,
request_conversion_transport_supported, request_conversion_transport_unsupported_reason,
request_pair_allowed_for_transport, CandidateTransportPolicyFacts,
request_pair_allowed_for_transport, request_pair_direct_auth,
request_pair_transport_unsupported_reason, CandidateTransportPolicyFacts,
};
pub(crate) use crate::control::GatewayControlDecision;
pub(crate) use crate::execution_runtime::{ConversionMode, ExecutionStrategy};
@@ -97,6 +98,28 @@ pub(crate) fn build_provider_transport_request_url(
)
}
pub(crate) fn build_provider_transport_request_url_for_request_body(
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
mapped_model: Option<&str>,
upstream_is_stream: bool,
request_query: Option<&str>,
kiro_api_region: Option<&str>,
provider_request_body: Option<&serde_json::Value>,
) -> Option<String> {
self::transport::build_transport_request_url_for_request_body(
transport,
self::transport::TransportRequestUrlParams {
provider_api_format,
mapped_model,
upstream_is_stream,
request_query,
kiro_api_region,
},
provider_request_body,
)
}
pub(crate) async fn resolve_execution_runtime_auth_context(
state: &AppState,
decision: &GatewayControlDecision,

5
apps/aether-gateway/src/ai_serving/planner/common.rs
View File

@@ -13,8 +13,9 @@ pub(crate) use crate::ai_serving::{
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_CHAT_STREAM_PLAN_KIND,
GEMINI_CHAT_SYNC_PLAN_KIND, GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND,
GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND,
GEMINI_FILES_LIST_PLAN_KIND, GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
GEMINI_EMBEDDING_SYNC_PLAN_KIND, GEMINI_FILES_DELETE_PLAN_KIND,
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND,
GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND,

24
apps/aether-gateway/src/ai_serving/planner/decision/control_plan.rs
View File

@@ -1,16 +1,16 @@
use crate::ai_serving::planner::common::{
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_FILES_DELETE_PLAN_KIND,
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND,
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND,
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND,
GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND,
GEMINI_FILES_LIST_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
OPENAI_RESPONSES_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
};
use crate::ai_serving::planner::plan_builders::{
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
@@ -110,7 +110,9 @@ fn build_sync_plan_payload_from_decision(
| OPENAI_RERANK_SYNC_PLAN_KIND => {
build_standard_sync_plan_from_decision(parts, body_json, payload)?
}
GEMINI_CHAT_SYNC_PLAN_KIND | GEMINI_CLI_SYNC_PLAN_KIND => {
GEMINI_CHAT_SYNC_PLAN_KIND
| GEMINI_CLI_SYNC_PLAN_KIND
| GEMINI_EMBEDDING_SYNC_PLAN_KIND => {
build_gemini_sync_plan_from_decision(parts, body_json, payload)?
}
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND

1
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs
View File

@@ -254,6 +254,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
spec,
prepared.upstream_is_stream,
prepared.kiro_auth.as_ref(),
Some(&provider_request_body),
) else {
mark_skipped_local_same_format_provider_candidate_with_failure_diagnostic(
state,

2
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/request/url.rs
View File

@@ -14,6 +14,7 @@ pub(crate) fn build_same_format_upstream_url(
spec: LocalSameFormatProviderSpec,
upstream_is_stream: bool,
kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>,
provider_request_body: Option<&serde_json::Value>,
) -> Option<String> {
build_same_format_provider_upstream_url_impl(
transport,
@@ -23,6 +24,7 @@ pub(crate) fn build_same_format_upstream_url(
upstream_is_stream,
request_query: parts.uri.query(),
kiro_api_region: kiro_auth.map(|auth| auth.auth_config.effective_api_region()),
provider_request_body,
},
)
}

1
apps/aether-gateway/src/ai_serving/planner/specialized/image/request.rs
View File

@@ -361,6 +361,7 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
&converted.mapped_model,
provider_api_format,
upstream_is_stream,
Some(&converted.body_json),
) else {
mark_skipped_local_openai_image_candidate_with_failure_diagnostic(
state,

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

@@ -75,15 +75,18 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
.await;
}
let is_kiro_claude_cli = is_kiro_claude_messages_transport(transport, provider_api_format);
let Some(conversion_kind) =
crate::ai_serving::request_conversion_kind(spec_metadata.api_format, provider_api_format)
else {
return None;
};
if let Some(skip_reason) = crate::ai_serving::request_conversion_transport_unsupported_reason(
if !crate::ai_serving::request_pair_allowed_for_transport(
transport,
conversion_kind,
spec_metadata.api_format,
provider_api_format,
) {
return None;
}
if let Some(skip_reason) = crate::ai_serving::request_pair_transport_unsupported_reason(
transport,
spec_metadata.api_format,
provider_api_format,
) {
mark_skipped_local_standard_candidate(
state,
@@ -156,7 +159,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
planner_state,
transport,
candidate,
crate::ai_serving::request_conversion_direct_auth(transport, conversion_kind),
crate::ai_serving::request_pair_direct_auth(transport, provider_api_format),
oauth_context,
)
.await
@@ -286,6 +289,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
&prepared_candidate.mapped_model,
provider_api_format,
upstream_is_stream,
Some(&provider_request_body),
) {
Some(url) => url,
None => {

13
apps/aether-gateway/src/ai_serving/planner/standard/mod.rs
View File

@@ -38,6 +38,7 @@ pub(crate) use self::openai::{
map_openai_reasoning_effort_to_gemini_budget, maybe_build_stream_local_decision_payload,
maybe_build_stream_local_openai_responses_decision_payload,
maybe_build_sync_local_decision_payload,
maybe_build_sync_local_openai_embedding_decision_payload,
maybe_build_sync_local_openai_responses_decision_payload, parse_openai_stop_sequences,
resolve_openai_chat_max_tokens, set_local_openai_chat_execution_exhausted_diagnostic,
value_as_u64,
@@ -66,14 +67,16 @@ pub(crate) fn build_standard_upstream_url(
mapped_model: &str,
provider_api_format: &str,
upstream_is_stream: bool,
provider_request_body: Option<&serde_json::Value>,
) -> Option<String> {
crate::ai_serving::build_provider_transport_request_url(
crate::ai_serving::build_provider_transport_request_url_for_request_body(
transport,
provider_api_format,
Some(mapped_model),
upstream_is_stream,
parts.uri.query(),
None,
provider_request_body,
)
}
@@ -85,6 +88,14 @@ pub(crate) async fn maybe_build_sync_local_standard_decision_payload(
body_json: &serde_json::Value,
plan_kind: &str,
) -> Result<Option<AiExecutionDecision>, GatewayError> {
if let Some(payload) = self::openai::maybe_build_sync_local_openai_embedding_decision_payload(
state, parts, trace_id, decision, body_json, plan_kind,
)
.await?
{
return Ok(Some(payload));
}
if let Some(payload) = self::claude::maybe_build_sync_local_claude_decision_payload(
state, parts, trace_id, decision, body_json, plan_kind,
)

31
apps/aether-gateway/src/ai_serving/planner/standard/openai/embedding.rs Normal file
View File

@@ -0,0 +1,31 @@
use crate::ai_serving::resolve_openai_embedding_sync_spec as resolve_surface_sync_spec;
use crate::ai_serving::GatewayControlDecision;
use crate::{AiExecutionDecision, AppState, GatewayError};
use super::super::family::maybe_build_sync_via_standard_family_payload;
pub(crate) fn resolve_sync_spec(
plan_kind: &str,
) -> Option<super::super::family::LocalStandardSpec> {
resolve_surface_sync_spec(plan_kind)
}
pub(crate) async fn maybe_build_sync_local_openai_embedding_decision_payload(
state: &AppState,
parts: &http::request::Parts,
trace_id: &str,
decision: &GatewayControlDecision,
body_json: &serde_json::Value,
plan_kind: &str,
) -> Result<Option<AiExecutionDecision>, GatewayError> {
maybe_build_sync_via_standard_family_payload(
state,
parts,
trace_id,
decision,
body_json,
plan_kind,
resolve_sync_spec,
)
.await
}

2
apps/aether-gateway/src/ai_serving/planner/standard/openai/mod.rs
View File

@@ -1,4 +1,5 @@
mod chat;
mod embedding;
mod responses;
pub(crate) use crate::ai_serving::{
@@ -14,6 +15,7 @@ pub(crate) use chat::{
maybe_build_stream_local_decision_payload, maybe_build_sync_local_decision_payload,
set_local_openai_chat_execution_exhausted_diagnostic,
};
pub(crate) use embedding::maybe_build_sync_local_openai_embedding_decision_payload;
pub(crate) use responses::{
build_local_openai_responses_stream_attempt_source_for_kind,
build_local_openai_responses_stream_plan_and_reports_for_kind,

5
apps/aether-gateway/src/ai_serving/pure/mod.rs
View File

@@ -80,8 +80,8 @@ pub(crate) use aether_ai_formats::api::{
resolve_local_image_stream_spec, resolve_local_image_sync_spec,
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
resolve_local_video_sync_spec, resolve_openai_chat_max_tokens,
resolve_openai_responses_stream_spec, resolve_openai_responses_sync_spec,
resolve_requested_gemini_image_model_for_request,
resolve_openai_embedding_sync_spec, resolve_openai_responses_stream_spec,
resolve_openai_responses_sync_spec, resolve_requested_gemini_image_model_for_request,
resolve_requested_openai_image_model_for_request,
resolve_upstream_is_stream_from_endpoint_config, sanitize_request_path,
sanitize_request_path_and_query, sanitize_request_query_string,
@@ -119,6 +119,7 @@ pub(crate) use aether_ai_formats::api::{
GEMINI_CLI_STREAM_SUCCESS_REPORT_KIND, GEMINI_CLI_SYNC_ERROR_REPORT_KIND,
GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND, GEMINI_CLI_SYNC_PLAN_KIND,
GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME,
GEMINI_EMBEDDING_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND,
GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND,
GEMINI_FILES_LIST_PLAN_KIND, GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
GEMINI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND, GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND,

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

@@ -61,18 +61,19 @@ pub(crate) use aether_provider_transport::{
build_same_format_provider_upstream_url, build_standard_plan_fallback_headers,
build_standard_plan_fallback_openai_chat_url,
build_standard_plan_fallback_openai_responses_url, build_standard_provider_request_headers,
build_transport_request_url, build_video_create_headers, build_video_create_request_body,
build_video_create_upstream_url, candidate_common_transport_skip_reason,
candidate_transport_pair_skip_reason, classify_same_format_provider_request_behavior,
ensure_upstream_auth_header, gemini_files_transport_unsupported_reason,
header_rules_are_locally_supported, header_rules_have_enabled_rules,
local_gemini_transport_unsupported_reason_with_network,
build_transport_request_url, build_transport_request_url_for_request_body,
build_video_create_headers, build_video_create_request_body, build_video_create_upstream_url,
candidate_common_transport_skip_reason, candidate_transport_pair_skip_reason,
classify_same_format_provider_request_behavior, ensure_upstream_auth_header,
gemini_files_transport_unsupported_reason, header_rules_are_locally_supported,
header_rules_have_enabled_rules, local_gemini_transport_unsupported_reason_with_network,
local_openai_chat_transport_unsupported_reason,
local_standard_transport_unsupported_reason_with_network,
openai_image_transport_unsupported_reason, request_conversion_direct_auth,
request_conversion_enabled_for_transport, request_conversion_transport_supported,
request_conversion_transport_unsupported_reason, request_pair_allowed_for_transport,
resolve_gemini_files_auth, resolve_openai_image_auth, resolve_same_format_provider_direct_auth,
request_pair_direct_auth, request_pair_transport_unsupported_reason, resolve_gemini_files_auth,
resolve_openai_image_auth, resolve_same_format_provider_direct_auth,
resolve_transport_execution_timeouts, resolve_transport_profile,
resolve_transport_proxy_snapshot, resolve_transport_proxy_snapshot_with_tunnel_affinity,
resolve_video_create_auth, same_format_provider_transport_supported,

11
apps/aether-gateway/src/control/route/ai.rs
View File

@@ -104,6 +104,17 @@ pub(super) fn classify_ai_public_route(
"gemini:video",
true,
))
} else if normalized_path.ends_with(":embedContent")
|| normalized_path.ends_with(":batchEmbedContents")
{
Some(classified_with_request_auth_channel(
"ai_public",
"gemini",
"embedding",
"api_key",
"gemini:embedding",
true,
))
} else if is_gemini_cli_request(headers) {
Some(classified_with_request_auth_channel(
"ai_public",

2
apps/aether-gateway/src/control/route/mod.rs
View File

@@ -229,6 +229,8 @@ pub(super) fn is_gemini_models_route(path: &str) -> bool {
(path.starts_with("/v1/models/") || path.starts_with("/v1beta/models/"))
&& (path.contains(":generateContent")
|| path.contains(":streamGenerateContent")
|| path.contains(":embedContent")
|| path.contains(":batchEmbedContents")
|| path.contains(":predictLongRunning"))
}

38
apps/aether-gateway/src/control/tests/ai.rs
View File

@@ -197,6 +197,44 @@ fn classifies_gemini_generate_content_api_key_without_cli_marker() {
assert!(decision.is_execution_runtime_candidate());
}
#[test]
fn classifies_gemini_embed_content_as_embedding_route() {
let headers = headers(&[("x-goog-api-key", "gemini-key")]);
let uri: Uri = "/v1beta/models/gemini-embedding-2-preview:embedContent"
.parse()
.expect("uri should parse");
let decision =
classify_control_route(&http::Method::POST, &uri, &headers).expect("route should classify");
assert_eq!(decision.route_family.as_deref(), Some("gemini"));
assert_eq!(decision.route_kind.as_deref(), Some("embedding"));
assert_eq!(decision.request_auth_channel.as_deref(), Some("api_key"));
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("gemini:embedding")
);
assert!(decision.is_execution_runtime_candidate());
}
#[test]
fn classifies_gemini_batch_embed_contents_as_embedding_route() {
let headers = headers(&[("x-goog-api-key", "gemini-key")]);
let uri: Uri = "/v1beta/models/gemini-embedding-2-preview:batchEmbedContents"
.parse()
.expect("uri should parse");
let decision =
classify_control_route(&http::Method::POST, &uri, &headers).expect("route should classify");
assert_eq!(decision.route_family.as_deref(), Some("gemini"));
assert_eq!(decision.route_kind.as_deref(), Some("embedding"));
assert_eq!(decision.request_auth_channel.as_deref(), Some("api_key"));
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("gemini:embedding")
);
assert!(decision.is_execution_runtime_candidate());
}
#[test]
fn classifies_gemini_predict_long_running_as_video_route() {
let headers = headers(&[]);

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

@@ -78,6 +78,87 @@ fn embedding_execution_runtime() -> Router {
)
}
fn gemini_embedding_success_state(
execution_runtime_url: String,
client_api_format: &str,
) -> AppState {
let mut snapshot = sample_currently_usable_auth_snapshot(
"key-gemini-embedding-success",
"user-gemini-embedding-success",
);
snapshot.user_allowed_providers = None;
snapshot.api_key_allowed_providers = None;
snapshot.user_allowed_api_formats = Some(vec![client_api_format.to_string()]);
snapshot.api_key_allowed_api_formats = Some(vec![client_api_format.to_string()]);
snapshot.user_allowed_models = Some(vec!["gemini-embedding-2-preview".to_string()]);
snapshot.api_key_allowed_models = Some(vec!["gemini-embedding-2-preview".to_string()]);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-gemini-embedding-success")),
snapshot,
)]));
let candidate_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
gemini_embedding_candidate_row(),
]));
let mut provider = sample_provider("provider-gemini-embedding", "Gemini Embeddings", 1);
provider.provider_type = "gemini".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-gemini-embedding",
"provider-gemini-embedding",
"gemini:embedding",
"https://generativelanguage.googleapis.com/v1beta",
)],
vec![sample_key(
"key-upstream-gemini-embedding",
"provider-gemini-embedding",
"gemini:embedding",
"sk-upstream-gemini-embedding",
)],
));
let data_state =
GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
provider_catalog_repository,
candidate_repository,
)
.with_auth_api_key_reader(auth_repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY);
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(data_state)
}
fn gemini_embedding_conversion_execution_runtime() -> Router {
Router::new().route(
"/v1/execute/sync",
any(|Json(plan): Json<ExecutionPlan>| async move {
assert_openai_to_gemini_embedding_execution_plan(&plan);
Json(gemini_embedding_execution_result(&plan))
}),
)
}
fn gemini_embedding_batch_conversion_execution_runtime() -> Router {
Router::new().route(
"/v1/execute/sync",
any(|Json(plan): Json<ExecutionPlan>| async move {
assert_openai_to_gemini_batch_embedding_execution_plan(&plan);
Json(gemini_batch_embedding_execution_result(&plan))
}),
)
}
fn gemini_embedding_native_execution_runtime() -> Router {
Router::new().route(
"/v1/execute/sync",
any(|Json(plan): Json<ExecutionPlan>| async move {
assert_native_gemini_embedding_execution_plan(&plan);
Json(gemini_embedding_execution_result(&plan))
}),
)
}
fn embedding_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-embedding".to_string(),
@@ -112,6 +193,40 @@ fn embedding_candidate_row() -> StoredMinimalCandidateSelectionRow {
}
}
fn gemini_embedding_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-gemini-embedding".to_string(),
provider_name: "Gemini Embeddings".to_string(),
provider_type: "gemini".to_string(),
provider_priority: 1,
provider_is_active: true,
endpoint_id: "endpoint-gemini-embedding".to_string(),
endpoint_api_format: "gemini:embedding".to_string(),
endpoint_api_family: Some("gemini".to_string()),
endpoint_kind: Some("embedding".to_string()),
endpoint_is_active: true,
key_id: "key-upstream-gemini-embedding".to_string(),
key_name: "default".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["gemini:embedding".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 50,
key_global_priority_by_format: None,
model_id: "model-gemini-embedding-preview".to_string(),
global_model_id: "global-gemini-embedding-preview".to_string(),
global_model_name: "gemini-embedding-2-preview".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(false),
model_provider_model_name: "gemini-embedding-2-preview".to_string(),
model_provider_model_mappings: None,
model_supports_streaming: Some(false),
model_is_active: true,
model_is_available: true,
}
}
fn assert_embedding_execution_plan(plan: &ExecutionPlan) {
assert_eq!(plan.client_api_format, "openai:embedding");
assert_eq!(plan.provider_api_format, "openai:embedding");
@@ -123,6 +238,78 @@ fn assert_embedding_execution_plan(plan: &ExecutionPlan) {
assert!(body.get("input").is_some());
}
fn assert_openai_to_gemini_embedding_execution_plan(plan: &ExecutionPlan) {
assert_eq!(plan.client_api_format, "openai:embedding");
assert_eq!(plan.provider_api_format, "gemini:embedding");
assert_eq!(plan.method, "POST");
assert_eq!(
plan.url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:embedContent"
);
assert_eq!(
plan.headers.get("x-goog-api-key").map(String::as_str),
Some("sk-upstream-gemini-embedding")
);
assert_eq!(
plan.model_name.as_deref(),
Some("gemini-embedding-2-preview")
);
assert!(!plan.stream);
let body = plan.body.json_body.as_ref().expect("json request body");
assert_eq!(body["model"], "gemini-embedding-2-preview");
assert_eq!(body["content"]["parts"][0]["text"], "hello");
assert!(body.get("input").is_none());
assert!(body.get("messages").is_none());
}
fn assert_openai_to_gemini_batch_embedding_execution_plan(plan: &ExecutionPlan) {
assert_eq!(plan.client_api_format, "openai:embedding");
assert_eq!(plan.provider_api_format, "gemini:embedding");
assert_eq!(plan.method, "POST");
assert_eq!(
plan.url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:batchEmbedContents"
);
assert_eq!(
plan.headers.get("x-goog-api-key").map(String::as_str),
Some("sk-upstream-gemini-embedding")
);
assert!(!plan.stream);
let body = plan.body.json_body.as_ref().expect("json request body");
assert!(body.get("model").is_none());
let requests = body["requests"].as_array().expect("batch requests");
assert_eq!(requests.len(), 2);
assert_eq!(requests[0]["model"], "models/gemini-embedding-2-preview");
assert_eq!(requests[0]["content"]["parts"][0]["text"], "hello");
assert_eq!(requests[1]["model"], "models/gemini-embedding-2-preview");
assert_eq!(requests[1]["content"]["parts"][0]["text"], "world");
assert!(body.get("input").is_none());
assert!(body.get("messages").is_none());
}
fn assert_native_gemini_embedding_execution_plan(plan: &ExecutionPlan) {
assert_eq!(plan.client_api_format, "gemini:embedding");
assert_eq!(plan.provider_api_format, "gemini:embedding");
assert_eq!(plan.method, "POST");
assert_eq!(
plan.url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:embedContent"
);
assert_eq!(
plan.headers.get("x-goog-api-key").map(String::as_str),
Some("sk-upstream-gemini-embedding")
);
assert_eq!(
plan.model_name.as_deref(),
Some("gemini-embedding-2-preview")
);
assert!(!plan.stream);
let body = plan.body.json_body.as_ref().expect("json request body");
assert_eq!(body["content"]["parts"][0]["text"], "hello");
assert!(body.get("input").is_none());
assert!(body.get("messages").is_none());
}
fn embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
ExecutionResult {
request_id: plan.request_id.clone(),
@@ -145,6 +332,55 @@ fn embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
}
}
fn 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!({
"model": "gemini-embedding-2-preview",
"embedding": {
"values": [0.1, 0.2, 0.3]
},
"usageMetadata": {
"promptTokenCount": 4,
"totalTokenCount": 4
}
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
}
}
fn gemini_batch_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!({
"model": "gemini-embedding-2-preview",
"embeddings": [
{"values": [0.1, 0.2, 0.3]},
{"values": [0.4, 0.5, 0.6]}
],
"usageMetadata": {
"promptTokenCount": 8,
"totalTokenCount": 8
}
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
}
}
#[tokio::test]
async fn embeddings_route_accepts_openai_payload() {
let (execution_runtime_url, execution_runtime_handle) =
@@ -215,6 +451,157 @@ async fn embeddings_route_accepts_openai_payload() {
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_converts_openai_payload_to_gemini_embedding_provider() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(gemini_embedding_conversion_execution_runtime()).await;
let gateway = build_router_with_state(gemini_embedding_success_state(
execution_runtime_url,
"openai:embedding",
));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/embeddings"))
.header(
http::header::AUTHORIZATION,
"Bearer sk-gemini-embedding-success",
)
.json(&json!({
"model": "gemini-embedding-2-preview",
"input": "hello"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
.and_then(|value| value.to_str().ok()),
Some("openai:embedding")
);
assert_eq!(
response
.headers()
.get(CONTROL_EXECUTION_RUNTIME_HEADER)
.and_then(|value| value.to_str().ok()),
Some("true")
);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(payload["object"], "list");
assert_eq!(payload["model"], "gemini-embedding-2-preview");
assert_eq!(payload["data"][0]["object"], "embedding");
assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3]));
assert_eq!(payload["usage"]["prompt_tokens"], json!(4));
assert_eq!(payload["usage"]["total_tokens"], json!(4));
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_converts_openai_batch_payload_to_gemini_batch_endpoint() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(gemini_embedding_batch_conversion_execution_runtime()).await;
let gateway = build_router_with_state(gemini_embedding_success_state(
execution_runtime_url,
"openai:embedding",
));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/embeddings"))
.header(
http::header::AUTHORIZATION,
"Bearer sk-gemini-embedding-success",
)
.json(&json!({
"model": "gemini-embedding-2-preview",
"input": ["hello", "world"]
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
.and_then(|value| value.to_str().ok()),
Some("openai:embedding")
);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(payload["object"], "list");
assert_eq!(payload["data"].as_array().map(Vec::len), Some(2));
assert_eq!(payload["data"][0]["index"], json!(0));
assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3]));
assert_eq!(payload["data"][1]["index"], json!(1));
assert_eq!(payload["data"][1]["embedding"], json!([0.4, 0.5, 0.6]));
assert_eq!(payload["usage"]["prompt_tokens"], json!(8));
assert_eq!(payload["usage"]["total_tokens"], json!(8));
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gemini_embed_content_route_uses_native_gemini_embedding_provider() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(gemini_embedding_native_execution_runtime()).await;
let gateway = build_router_with_state(gemini_embedding_success_state(
execution_runtime_url,
"gemini:embedding",
));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/v1beta/models/gemini-embedding-2-preview:embedContent"
))
.header("x-goog-api-key", "sk-gemini-embedding-success")
.json(&json!({
"content": {
"parts": [{"text": "hello"}]
}
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_FAMILY_HEADER)
.and_then(|value| value.to_str().ok()),
Some("gemini")
);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_KIND_HEADER)
.and_then(|value| value.to_str().ok()),
Some("embedding")
);
assert_eq!(
response
.headers()
.get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
.and_then(|value| value.to_str().ok()),
Some("gemini:embedding")
);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(payload["embedding"]["values"], json!([0.1, 0.2, 0.3]));
assert_eq!(payload["model"], "gemini-embedding-2-preview");
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_accepts_all_canonical_input_shapes() {
let (execution_runtime_url, execution_runtime_handle) =