mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
@@ -50,13 +50,14 @@ 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,
|
||||
sanitize_request_path_and_query, AiControlPlanRequest, CanonicalContentPart,
|
||||
CanonicalStreamEvent, CanonicalStreamFrame, ClaudeClientEmitter, ExecutionRuntimeAuthContext,
|
||||
LocalCoreSyncErrorKind, LocalOpenAiImageSpec, LocalSameFormatProviderFamily,
|
||||
LocalSameFormatProviderSpec, LocalStandardSourceFamily, LocalStandardSourceMode,
|
||||
LocalStandardSpec, OpenAIChatClientEmitter, OpenAIResponsesClientEmitter,
|
||||
StreamingStandardTerminalObserver, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_FILES_DOWNLOAD_PLAN_KIND,
|
||||
resolve_openai_embedding_sync_spec, sanitize_request_path_and_query, AiControlPlanRequest,
|
||||
CanonicalContentPart, CanonicalStreamEvent, CanonicalStreamFrame, ClaudeClientEmitter,
|
||||
ExecutionRuntimeAuthContext, LocalCoreSyncErrorKind, LocalOpenAiImageSpec,
|
||||
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, LocalStandardSourceFamily,
|
||||
LocalStandardSourceMode, LocalStandardSpec, OpenAIChatClientEmitter,
|
||||
OpenAIResponsesClientEmitter, 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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -258,7 +258,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
|
||||
.eq_ignore_ascii_case("grok");
|
||||
let transport_profile =
|
||||
crate::ai_serving::transport::resolve_transport_profile(&prepared.transport);
|
||||
let Some(upstream_url) = (if is_grok {
|
||||
let upstream_url = if is_grok {
|
||||
Some(build_grok_upstream_url(&prepared.transport, GROK_CHAT_PATH))
|
||||
} else {
|
||||
super::super::request::build_same_format_upstream_url(
|
||||
@@ -269,8 +269,10 @@ 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 {
|
||||
};
|
||||
let Some(upstream_url) = upstream_url else {
|
||||
mark_skipped_local_same_format_provider_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
input,
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
@@ -393,6 +393,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,
|
||||
|
||||
@@ -190,15 +190,18 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
|
||||
});
|
||||
}
|
||||
|
||||
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,
|
||||
@@ -271,7 +274,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
|
||||
@@ -353,6 +356,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),
|
||||
);
|
||||
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,
|
||||
@@ -373,6 +399,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),
|
||||
);
|
||||
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() {
|
||||
@@ -401,6 +450,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 => {
|
||||
@@ -480,6 +530,18 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
|
||||
})
|
||||
}
|
||||
|
||||
fn apply_transport_request_body_semantics(
|
||||
provider_request_body: &mut Value,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
provider_api_format: &str,
|
||||
) -> 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(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
|
||||
@@ -40,6 +40,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,
|
||||
@@ -68,14 +69,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,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -87,6 +90,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,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -1,5 +1,6 @@
|
||||
mod chat;
|
||||
mod image_intent;
|
||||
mod embedding;
|
||||
mod responses;
|
||||
|
||||
pub(crate) use crate::ai_serving::{
|
||||
@@ -16,6 +17,7 @@ pub(crate) use chat::{
|
||||
set_local_openai_chat_execution_exhausted_diagnostic,
|
||||
};
|
||||
pub(super) use image_intent::openai_request_is_image_generation_intent;
|
||||
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,
|
||||
|
||||
@@ -81,8 +81,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,
|
||||
@@ -120,6 +120,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,
|
||||
|
||||
@@ -55,7 +55,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_grok_app_chat_body, build_grok_browser_headers, build_grok_upstream_url,
|
||||
@@ -66,24 +67,24 @@ 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_grok_session_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, same_format_provider_transport_unsupported_reason,
|
||||
should_skip_upstream_passthrough_header, should_try_same_format_provider_oauth_auth,
|
||||
supports_local_gemini_transport_with_network,
|
||||
request_pair_direct_auth, request_pair_transport_unsupported_reason, resolve_gemini_files_auth,
|
||||
resolve_grok_session_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,
|
||||
same_format_provider_transport_unsupported_reason, should_skip_upstream_passthrough_header,
|
||||
should_try_same_format_provider_oauth_auth, supports_local_gemini_transport_with_network,
|
||||
supports_local_generic_oauth_request_auth_resolution,
|
||||
supports_local_oauth_request_auth_resolution, transport_proxy_is_locally_supported,
|
||||
video_create_transport_unsupported_reason, CandidateTransportPolicyFacts,
|
||||
@@ -96,4 +97,5 @@ pub(crate) use aether_provider_transport::{
|
||||
StandardPlanFallbackHeadersInput, StandardProviderRequestHeaders,
|
||||
StandardProviderRequestHeadersInput, TransportRequestUrlParams, GROK_CHAT_PATH,
|
||||
GROK_INTERNAL_HEADER, GROK_RATE_LIMITS_PATH,
|
||||
TransportRequestBodySemanticsError,
|
||||
};
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -85,7 +85,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"),
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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"))
|
||||
}
|
||||
|
||||
|
||||
@@ -219,6 +219,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(&[]);
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -62,6 +63,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,
|
||||
@@ -344,14 +354,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,
|
||||
@@ -625,6 +635,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 {
|
||||
|
||||
@@ -12,7 +12,7 @@ use crate::control::GatewayControlDecision;
|
||||
use crate::usage::spawn_sync_report;
|
||||
use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError};
|
||||
use axum::body::Body;
|
||||
use axum::http::Response;
|
||||
use axum::http::{Response, StatusCode};
|
||||
use base64::Engine as _;
|
||||
use tracing::warn;
|
||||
|
||||
@@ -148,6 +148,74 @@ fn build_local_core_sync_finalize_fallback_response(
|
||||
build_local_sync_response_from_bytes(trace_id, decision, payload, Vec::new())
|
||||
}
|
||||
|
||||
fn maybe_build_invalid_provider_success_finalize_response(
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
if !local_core_sync_finalize_has_invalid_provider_success(payload)? {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let client_api_format = resolve_local_sync_client_api_format(payload);
|
||||
let message = "Provider returned HTTP 200 but the Gemini response did not contain visible model output; refusing to finalize it as a successful response.";
|
||||
let body_json = build_core_error_body_for_client_format(
|
||||
&client_api_format,
|
||||
message,
|
||||
Some("invalid_provider_success_response"),
|
||||
LocalCoreSyncErrorKind::ServerError,
|
||||
)
|
||||
.unwrap_or_else(|| {
|
||||
serde_json::json!({
|
||||
"error": {
|
||||
"message": message,
|
||||
"type": "server_error",
|
||||
"code": "invalid_provider_success_response"
|
||||
}
|
||||
})
|
||||
});
|
||||
|
||||
let mut response_headers = payload.headers.clone();
|
||||
response_headers.remove("content-encoding");
|
||||
response_headers.remove("content-length");
|
||||
response_headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
let body_bytes =
|
||||
serde_json::to_vec(&body_json).map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
response_headers.insert("content-length".to_string(), body_bytes.len().to_string());
|
||||
|
||||
Ok(Some(build_client_response_from_parts(
|
||||
StatusCode::BAD_GATEWAY.as_u16(),
|
||||
&response_headers,
|
||||
Body::from(body_bytes),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)?))
|
||||
}
|
||||
|
||||
fn local_core_sync_finalize_has_invalid_provider_success(
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Result<bool, GatewayError> {
|
||||
if payload.status_code >= 400 || !is_core_error_finalize_kind(payload.report_kind.as_str()) {
|
||||
return Ok(false);
|
||||
}
|
||||
let provider_api_format = resolve_local_sync_provider_api_format(payload);
|
||||
if aether_ai_formats::normalize_api_format_alias(&provider_api_format)
|
||||
!= "gemini:generate_content"
|
||||
{
|
||||
return Ok(false);
|
||||
}
|
||||
let Some(body_json) = resolve_local_sync_source_body_json(payload)? else {
|
||||
return Ok(false);
|
||||
};
|
||||
if has_nested_error(&body_json) {
|
||||
return Ok(false);
|
||||
}
|
||||
Ok(
|
||||
aether_ai_formats::formats::gemini::generate_content::response::from_raw(&body_json)
|
||||
.is_none(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_best_effort_local_core_error_body(
|
||||
payload: &GatewaySyncReportRequest,
|
||||
body_json: &serde_json::Value,
|
||||
@@ -283,6 +351,16 @@ fn resolve_local_sync_client_api_format(payload: &GatewaySyncReportRequest) -> S
|
||||
.to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn resolve_local_sync_provider_api_format(payload: &GatewaySyncReportRequest) -> String {
|
||||
payload
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("provider_api_format"))
|
||||
.and_then(|value| value.as_str())
|
||||
.map(|value| value.trim().to_ascii_lowercase())
|
||||
.unwrap_or_else(|| resolve_local_sync_client_api_format(payload))
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_core_error_background_report_kind(report_kind: &str) -> Option<String> {
|
||||
core_error_background_report_kind(report_kind).map(ToOwned::to_owned)
|
||||
}
|
||||
@@ -522,6 +600,10 @@ pub(crate) async fn submit_local_core_error_or_sync_finalize(
|
||||
maybe_compile_sync_finalize_response(trace_id, decision, &payload)?
|
||||
{
|
||||
response
|
||||
} else if let Some(response) =
|
||||
maybe_build_invalid_provider_success_finalize_response(trace_id, decision, &payload)?
|
||||
{
|
||||
response
|
||||
} else if let Some(response) =
|
||||
maybe_build_local_core_error_response(trace_id, decision, &payload)?
|
||||
{
|
||||
@@ -566,9 +648,10 @@ mod tests {
|
||||
use axum::body::to_bytes;
|
||||
use serde_json::json;
|
||||
|
||||
use super::maybe_build_local_core_error_response;
|
||||
use super::{maybe_build_local_core_error_response, submit_local_core_error_or_sync_finalize};
|
||||
use crate::control::GatewayControlDecision;
|
||||
use crate::usage::GatewaySyncReportRequest;
|
||||
use crate::AppState;
|
||||
|
||||
fn test_decision() -> GatewayControlDecision {
|
||||
GatewayControlDecision::synthetic(
|
||||
@@ -684,4 +767,59 @@ mod tests {
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_core_sync_finalize_rejects_gemini_http_200_without_visible_output() {
|
||||
let mut payload = core_finalize_payload(
|
||||
"openai_chat_sync_finalize",
|
||||
"openai:chat",
|
||||
"gemini:generate_content",
|
||||
200,
|
||||
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"
|
||||
}),
|
||||
);
|
||||
payload.report_context = Some(json!({
|
||||
"client_api_format": "openai:chat",
|
||||
"provider_api_format": "gemini:generate_content",
|
||||
"needs_conversion": true,
|
||||
"has_envelope": false
|
||||
}));
|
||||
|
||||
let state = AppState::new().expect("state should build");
|
||||
let response = submit_local_core_error_or_sync_finalize(
|
||||
&state,
|
||||
"trace-invalid-gemini-200",
|
||||
&test_decision(),
|
||||
payload,
|
||||
)
|
||||
.await
|
||||
.expect("response should build");
|
||||
|
||||
assert_eq!(response.status(), http::StatusCode::BAD_GATEWAY);
|
||||
let body: serde_json::Value = serde_json::from_slice(
|
||||
&to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read"),
|
||||
)
|
||||
.expect("body should decode");
|
||||
let message = body["error"]["message"]
|
||||
.as_str()
|
||||
.expect("error message should exist");
|
||||
assert!(
|
||||
message.contains("visible model output"),
|
||||
"unexpected message: {message}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,7 +3,10 @@ use std::io::Error as IoError;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use aether_contracts::{ExecutionPlan, ExecutionResult, ExecutionTelemetry};
|
||||
use aether_contracts::{
|
||||
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionResult,
|
||||
ExecutionTelemetry,
|
||||
};
|
||||
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
|
||||
use aether_scheduler_core::{
|
||||
execution_error_details, parse_request_candidate_report_context,
|
||||
@@ -26,8 +29,8 @@ use tokio::time::MissedTickBehavior;
|
||||
use tracing::{debug, warn};
|
||||
|
||||
use crate::ai_serving::api::{
|
||||
implicit_sync_finalize_report_kind, maybe_build_sync_finalize_outcome,
|
||||
LocalCoreSyncFinalizeOutcome,
|
||||
build_core_error_body_for_client_format, implicit_sync_finalize_report_kind,
|
||||
maybe_build_sync_finalize_outcome, LocalCoreSyncErrorKind, LocalCoreSyncFinalizeOutcome,
|
||||
};
|
||||
use crate::api::response::{
|
||||
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
|
||||
@@ -184,6 +187,55 @@ fn build_sync_report_payload(
|
||||
}
|
||||
}
|
||||
|
||||
fn invalid_gemini_provider_success_message(
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&Value>,
|
||||
status_code: u16,
|
||||
body_json: Option<&Value>,
|
||||
) -> Option<&'static str> {
|
||||
if status_code >= 400 {
|
||||
return None;
|
||||
}
|
||||
let provider_api_format = report_context
|
||||
.and_then(|value| value.get("provider_api_format"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or(plan.provider_api_format.as_str());
|
||||
if aether_ai_formats::normalize_api_format_alias(provider_api_format)
|
||||
!= "gemini:generate_content"
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let body_json = body_json?;
|
||||
if body_json
|
||||
.as_object()
|
||||
.is_some_and(|object| object.get("error").is_some_and(|error| !error.is_null()))
|
||||
{
|
||||
return None;
|
||||
}
|
||||
if aether_ai_formats::formats::gemini::generate_content::response::from_raw(body_json).is_some()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
Some("Provider returned HTTP 200 but the Gemini response did not contain visible model output; refusing to finalize it as a successful response.")
|
||||
}
|
||||
|
||||
fn build_invalid_provider_success_body(
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&Value>,
|
||||
message: &str,
|
||||
) -> Option<Value> {
|
||||
let client_api_format = report_context
|
||||
.and_then(|value| value.get("client_api_format"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or(plan.client_api_format.as_str());
|
||||
build_core_error_body_for_client_format(
|
||||
client_api_format,
|
||||
message,
|
||||
Some("invalid_provider_success_response"),
|
||||
LocalCoreSyncErrorKind::ServerError,
|
||||
)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct OpenAiImageSyncProgressSnapshot {
|
||||
phase: &'static str,
|
||||
@@ -1446,19 +1498,37 @@ async fn execute_execution_runtime_sync_impl(
|
||||
local_failover_response_text,
|
||||
local_failover_analysis,
|
||||
) = loop {
|
||||
let result_body_json = result
|
||||
.body
|
||||
.as_ref()
|
||||
.and_then(|body| body.json_body.as_ref());
|
||||
let (result_error_type, result_error_message) =
|
||||
execution_error_details(result.error.as_ref(), result_body_json);
|
||||
let result_latency_ms = result
|
||||
.telemetry
|
||||
.as_ref()
|
||||
.and_then(|telemetry| telemetry.elapsed_ms);
|
||||
let mut headers = std::mem::take(&mut result.headers);
|
||||
let (body_bytes, body_json, body_base64) =
|
||||
let (body_bytes, mut body_json, body_base64) =
|
||||
decode_execution_result_body(result.body.take(), &mut headers)?;
|
||||
if let Some(message) = invalid_gemini_provider_success_message(
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
result.status_code,
|
||||
body_json.as_ref(),
|
||||
) {
|
||||
result.status_code = StatusCode::BAD_GATEWAY.as_u16();
|
||||
result.error = Some(ExecutionError {
|
||||
kind: ExecutionErrorKind::Upstream5xx,
|
||||
phase: ExecutionPhase::Finalize,
|
||||
message: message.to_string(),
|
||||
upstream_status: Some(StatusCode::OK.as_u16()),
|
||||
retryable: false,
|
||||
failover_recommended: false,
|
||||
});
|
||||
if let Some(error_body) =
|
||||
build_invalid_provider_success_body(&plan, report_context.as_ref(), message)
|
||||
{
|
||||
body_json = Some(error_body);
|
||||
headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
}
|
||||
}
|
||||
let (result_error_type, result_error_message) =
|
||||
execution_error_details(result.error.as_ref(), body_json.as_ref());
|
||||
let local_failover_response_text = local_failover_response_text(
|
||||
body_json.as_ref(),
|
||||
&body_bytes,
|
||||
@@ -2167,6 +2237,41 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
fn test_gemini_chat_plan() -> ExecutionPlan {
|
||||
let mut plan = test_openai_image_plan(false);
|
||||
plan.client_api_format = "openai:chat".to_string();
|
||||
plan.provider_api_format = "gemini:generate_content".to_string();
|
||||
plan.model_name = Some("gemini-3-flash-preview".to_string());
|
||||
plan
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_gemini_provider_success_uses_plan_format_when_context_is_missing() {
|
||||
let plan = test_gemini_chat_plan();
|
||||
let body = json!({
|
||||
"candidates": [{
|
||||
"content": {"role": "model"},
|
||||
"finishReason": "MAX_TOKENS"
|
||||
}],
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 8,
|
||||
"candidatesTokenCount": 1,
|
||||
"thoughtsTokenCount": 25,
|
||||
"totalTokenCount": 34
|
||||
}
|
||||
});
|
||||
|
||||
let message = invalid_gemini_provider_success_message(
|
||||
&plan,
|
||||
None,
|
||||
StatusCode::OK.as_u16(),
|
||||
Some(&body),
|
||||
)
|
||||
.expect("empty Gemini 200 response should be rejected from plan format");
|
||||
|
||||
assert!(message.contains("visible model output"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn json_whitespace_heartbeat_stream_prefixes_final_json() {
|
||||
let (tx, rx) = mpsc::channel::<Result<Bytes, IoError>>(1);
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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 => {
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -1154,6 +1154,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"));
|
||||
|
||||
@@ -78,6 +78,147 @@ 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 vertex_gemini_embedding_success_state(execution_runtime_url: String) -> AppState {
|
||||
let mut snapshot = sample_currently_usable_auth_snapshot(
|
||||
"key-vertex-gemini-embedding-success",
|
||||
"user-vertex-gemini-embedding-success",
|
||||
);
|
||||
snapshot.user_allowed_providers = None;
|
||||
snapshot.api_key_allowed_providers = Some(vec!["openai".to_string(), "vertex_ai".to_string()]);
|
||||
snapshot.user_allowed_api_formats = Some(vec!["openai:embedding".to_string()]);
|
||||
snapshot.api_key_allowed_api_formats = Some(vec!["openai:embedding".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-vertex-gemini-embedding-success")),
|
||||
snapshot,
|
||||
)]));
|
||||
let candidate_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
vertex_gemini_embedding_candidate_row(),
|
||||
]));
|
||||
let mut provider = sample_provider("provider-vertex-gemini-embedding", "Vertex AI", 1);
|
||||
provider.provider_type = "vertex_ai".to_string();
|
||||
let mut key = sample_key(
|
||||
"key-upstream-vertex-gemini-embedding",
|
||||
"provider-vertex-gemini-embedding",
|
||||
"gemini:embedding",
|
||||
"sk-upstream-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-gemini-embedding",
|
||||
"gemini:embedding",
|
||||
"https://aiplatform.googleapis.com",
|
||||
)],
|
||||
vec![key],
|
||||
));
|
||||
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 vertex_gemini_embedding_conversion_execution_runtime() -> Router {
|
||||
Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(|Json(plan): Json<ExecutionPlan>| async move {
|
||||
assert_openai_to_vertex_gemini_embedding_execution_plan(&plan);
|
||||
Json(vertex_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 +253,53 @@ 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 vertex_gemini_embedding_candidate_row() -> StoredMinimalCandidateSelectionRow {
|
||||
let mut row = gemini_embedding_candidate_row();
|
||||
row.provider_id = "provider-vertex-gemini-embedding".to_string();
|
||||
row.provider_name = "Vertex AI".to_string();
|
||||
row.provider_type = "vertex_ai".to_string();
|
||||
row.endpoint_id = "endpoint-vertex-gemini-embedding".to_string();
|
||||
row.key_id = "key-upstream-vertex-gemini-embedding".to_string();
|
||||
row.key_name = "default".to_string();
|
||||
row.key_allowed_models = Some(vec!["gemini-embedding-2".to_string()]);
|
||||
row.model_provider_model_name = "gemini-embedding-2".to_string();
|
||||
row
|
||||
}
|
||||
|
||||
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 +311,103 @@ 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_vertex_gemini_embedding_execution_plan(plan: &ExecutionPlan) {
|
||||
assert_eq!(plan.provider_id, "provider-vertex-gemini-embedding");
|
||||
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://aiplatform.googleapis.com/v1/publishers/google/models/gemini-embedding-2:predict?key=sk-upstream-vertex-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!(
|
||||
body.get("model").is_none(),
|
||||
"Vertex predict carries the model in the path; the body must not repeat it"
|
||||
);
|
||||
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());
|
||||
}
|
||||
|
||||
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 +430,79 @@ 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 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(),
|
||||
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 +573,201 @@ 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_payload_to_vertex_gemini_embedding_provider() {
|
||||
let (execution_runtime_url, execution_runtime_handle) =
|
||||
start_server(vertex_gemini_embedding_conversion_execution_runtime()).await;
|
||||
let gateway =
|
||||
build_router_with_state(vertex_gemini_embedding_success_state(execution_runtime_url));
|
||||
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-vertex-gemini-embedding-success",
|
||||
)
|
||||
.json(&json!({
|
||||
"model": "gemini-embedding-2-preview",
|
||||
"input": "hello"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let endpoint_signature = response
|
||||
.headers()
|
||||
.get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_string);
|
||||
let status = response.status();
|
||||
let body_text = response.text().await.expect("body should read");
|
||||
assert_eq!(
|
||||
status,
|
||||
StatusCode::OK,
|
||||
"unexpected response body: {body_text}"
|
||||
);
|
||||
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");
|
||||
assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3]));
|
||||
|
||||
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) =
|
||||
|
||||
@@ -1951,6 +1951,10 @@ async fn gateway_handles_public_test_connection_without_hitting_fallback_probe()
|
||||
);
|
||||
assert_eq!(body_json["model"], "gpt-5");
|
||||
assert_eq!(body_json["messages"][0]["content"], "Health check");
|
||||
assert!(
|
||||
body_json.get("max_tokens").is_none(),
|
||||
"public OpenAI-compatible test connection must not force a tiny max_tokens value"
|
||||
);
|
||||
Json(json!({"id": "resp_local_test"})).into_response()
|
||||
}
|
||||
}),
|
||||
@@ -2009,6 +2013,91 @@ async fn gateway_handles_public_test_connection_without_hitting_fallback_probe()
|
||||
provider_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_gemini_test_connection_does_not_force_low_max_output_tokens() {
|
||||
let provider_hits = Arc::new(Mutex::new(0usize));
|
||||
let provider_hits_clone = Arc::clone(&provider_hits);
|
||||
let provider = Router::new().route(
|
||||
"/{*path}",
|
||||
any(move |request: Request| {
|
||||
let provider_hits_inner = Arc::clone(&provider_hits_clone);
|
||||
async move {
|
||||
*provider_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
let body = to_bytes(request.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read");
|
||||
let body_json: serde_json::Value =
|
||||
serde_json::from_slice(&body).expect("json body should parse");
|
||||
assert_eq!(body_json["contents"][0]["parts"][0]["text"], "Health check");
|
||||
assert!(
|
||||
body_json
|
||||
.get("generationConfig")
|
||||
.and_then(|config| config.get("maxOutputTokens"))
|
||||
.is_none(),
|
||||
"Gemini test connection must not force a tiny maxOutputTokens value"
|
||||
);
|
||||
Json(json!({
|
||||
"candidates": [{
|
||||
"content": {
|
||||
"role": "model",
|
||||
"parts": [{"text": "ok"}]
|
||||
},
|
||||
"finishReason": "STOP"
|
||||
}],
|
||||
"responseId": "gemini_test_connection_ok"
|
||||
}))
|
||||
.into_response()
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (provider_url, provider_handle) = start_server(provider).await;
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-gemini", "google", 10)],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-gemini",
|
||||
"provider-gemini",
|
||||
"gemini:generate_content",
|
||||
&provider_url,
|
||||
)],
|
||||
vec![sample_key(
|
||||
"key-gemini",
|
||||
"provider-gemini",
|
||||
"gemini:generate_content",
|
||||
"google-api-key",
|
||||
)],
|
||||
));
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/v1/test-connection?provider=provider-gemini&model=gemini-3-flash-preview&api_format=gemini:generate_content"
|
||||
))
|
||||
.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["status"], "success");
|
||||
assert_eq!(payload["provider_id"], "provider-gemini");
|
||||
assert_eq!(payload["endpoint_id"], "endpoint-gemini");
|
||||
assert_eq!(payload["api_format"], "gemini:generate_content");
|
||||
assert_eq!(*provider_hits.lock().expect("mutex should lock"), 1);
|
||||
|
||||
gateway_handle.abort();
|
||||
provider_handle.abort();
|
||||
}
|
||||
|
||||
async fn assert_public_support_route_returns_local_503(
|
||||
method: reqwest::Method,
|
||||
path: &str,
|
||||
|
||||
Reference in New Issue
Block a user