Merge pull request #371 from Kayphoon/feature/embedding-model-support

feat: add embedding and rerank support
This commit is contained in:
fawney19
2026-05-04 12:56:11 +08:00
committed by GitHub
87 changed files with 5517 additions and 184 deletions

View File

@@ -53,10 +53,11 @@ pub(crate) use aether_ai_formats::api::{
LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec,
StreamingStandardTerminalObserver, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_FILES_DOWNLOAD_PLAN_KIND,
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
};
pub(crate) fn parse_direct_request_body(

View File

@@ -13,12 +13,12 @@ pub(crate) use crate::ai_serving::{
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_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_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,
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,
};
pub(crate) use aether_ai_serving::AiRequestedModelFamily as RequestedModelFamily;

View File

@@ -4,12 +4,13 @@ use crate::ai_serving::planner::common::{
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_IMAGE_STREAM_PLAN_KIND,
OPENAI_IMAGE_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,
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,
@@ -103,7 +104,10 @@ fn build_sync_plan_payload_from_decision(
OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND => {
build_openai_responses_sync_plan_from_decision(parts, body_json, payload, true)?
}
CLAUDE_CHAT_SYNC_PLAN_KIND | CLAUDE_CLI_SYNC_PLAN_KIND => {
CLAUDE_CHAT_SYNC_PLAN_KIND
| CLAUDE_CLI_SYNC_PLAN_KIND
| OPENAI_EMBEDDING_SYNC_PLAN_KIND
| 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 => {

View File

@@ -84,6 +84,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
json!(super::super::ANTIGRAVITY_ENVELOPE_NAME),
);
}
let provider_api_format = resolved.provider_api_format.clone();
let report_context = append_local_failover_policy_to_value(
append_execution_contract_fields_to_value(
build_local_execution_report_context(LocalExecutionReportContextParts {
@@ -100,7 +101,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: spec_metadata.api_format,
provider_api_format: provider_api_format.as_str(),
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
@@ -126,7 +127,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
execution_strategy,
conversion_mode,
spec_metadata.api_format,
spec_metadata.api_format,
provider_api_format.as_str(),
),
&resolved.transport,
);
@@ -136,6 +137,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
is_kiro: _,
auth_header,
auth_value,
provider_api_format,
mapped_model,
report_kind,
upstream_is_stream,
@@ -161,7 +163,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
provider_request_method: None,
auth_header,
auth_value,
provider_api_format: spec_metadata.api_format.to_string(),
provider_api_format,
client_api_format: spec_metadata.api_format.to_string(),
model_name: input.requested_model.clone(),
mapped_model,

View File

@@ -29,12 +29,63 @@ use super::{
};
use crate::ai_serving::planner::standard::same_format_provider_request_body_failure_extra_data;
pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trace(
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
) -> Option<&'static str> {
let provider_api_format =
match crate::ai_serving::normalize_api_format_alias(provider_api_format).as_str() {
"openai:chat" => "openai:chat",
"openai:responses" => "openai:responses",
"openai:responses:compact" => "openai:responses:compact",
"openai:embedding" => "openai:embedding",
"openai:rerank" => "openai:rerank",
"claude:messages" => "claude:messages",
"gemini:generate_content" => "gemini:generate_content",
"gemini:embedding" => "gemini:embedding",
"jina:embedding" => "jina:embedding",
"jina:rerank" => "jina:rerank",
"doubao:embedding" => "doubao:embedding",
_ => return Some("transport_api_format_unsupported"),
};
let behavior = policy::classify_same_format_provider_request_behavior(
transport,
crate::ai_serving::planner::spec_metadata::LocalExecutionSurfaceSpecMetadata {
api_format: provider_api_format,
require_streaming: false,
requested_model_family: None,
decision_kind: "trace_candidate_metadata",
report_kind: Some("trace_candidate_metadata"),
},
);
if !behavior.is_antigravity
&& !behavior.is_claude_code
&& !behavior.is_vertex
&& !behavior.is_kiro
{
return None;
}
let family = if provider_api_format.starts_with("gemini:") {
crate::ai_serving::LocalSameFormatProviderFamily::Gemini
} else {
crate::ai_serving::LocalSameFormatProviderFamily::Standard
};
policy::same_format_provider_transport_unsupported_reason(
&behavior,
transport,
family,
provider_api_format,
)
}
pub(crate) struct LocalSameFormatProviderCandidatePayloadParts {
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
pub(super) is_antigravity: bool,
pub(super) is_kiro: bool,
pub(super) auth_header: Option<String>,
pub(super) auth_value: Option<String>,
pub(super) provider_api_format: String,
pub(super) mapped_model: String,
pub(super) report_kind: &'static str,
pub(super) upstream_is_stream: bool,
@@ -74,6 +125,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
let Some(mut base_provider_request_body) =
super::super::request::build_same_format_provider_request_body(
body_json,
prepared.provider_api_format.as_str(),
&prepared.mapped_model,
spec,
prepared.transport.endpoint.body_rules.as_ref(),
@@ -180,6 +232,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
parts,
&prepared.transport,
&prepared.mapped_model,
prepared.provider_api_format.as_str(),
spec,
prepared.upstream_is_stream,
prepared.kiro_auth.as_ref(),
@@ -248,6 +301,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
is_kiro: prepared.is_kiro,
auth_header: prepared.auth_header,
auth_value: prepared.auth_value,
provider_api_format: prepared.provider_api_format,
mapped_model: prepared.mapped_model,
report_kind: prepared.report_kind,
upstream_is_stream: prepared.upstream_is_stream,

View File

@@ -31,6 +31,7 @@ pub(super) struct PreparedSameFormatProviderCandidate {
pub(super) kiro_auth: Option<KiroRequestAuth>,
pub(super) auth_header: Option<String>,
pub(super) auth_value: Option<String>,
pub(super) provider_api_format: String,
pub(super) mapped_model: String,
pub(super) report_kind: &'static str,
pub(super) upstream_is_stream: bool,
@@ -49,19 +50,20 @@ pub(super) async fn prepare_local_same_format_provider_candidate(
let planner_state = PlannerAppState::new(state);
let candidate = &eligible.candidate;
let transport = Arc::clone(&eligible.transport);
let provider_api_format = eligible.provider_api_format.as_str();
let behavior = classify_same_format_provider_request_behavior(&transport, spec_metadata);
if !same_format_provider_transport_supported(
&behavior,
&transport,
spec.family,
spec_metadata.api_format,
provider_api_format,
) {
let skip_reason = same_format_provider_transport_unsupported_reason(
&behavior,
&transport,
spec.family,
spec_metadata.api_format,
provider_api_format,
)
.unwrap_or("transport_unsupported");
super::super::payload::mark_skipped_local_same_format_provider_candidate(
@@ -90,7 +92,7 @@ pub(super) async fn prepare_local_same_format_provider_candidate(
&transport,
OauthPreparationContext {
trace_id,
api_format: spec_metadata.api_format,
api_format: provider_api_format,
operation: "same_format_provider_prepare",
},
)
@@ -168,6 +170,7 @@ pub(super) async fn prepare_local_same_format_provider_candidate(
kiro_auth,
auth_header,
auth_value,
provider_api_format: provider_api_format.to_string(),
mapped_model,
report_kind: behavior.report_kind,
upstream_is_stream: behavior.upstream_is_stream,

View File

@@ -8,6 +8,7 @@ use crate::ai_serving::transport::{
pub(crate) fn build_same_format_provider_request_body(
body_json: &Value,
provider_api_format: &str,
mapped_model: &str,
spec: LocalSameFormatProviderSpec,
body_rules: Option<&Value>,
@@ -19,7 +20,8 @@ pub(crate) fn build_same_format_provider_request_body(
build_same_format_provider_request_body_impl(SameFormatProviderRequestBodyInput {
body_json,
mapped_model,
provider_api_format: spec.api_format,
client_api_format: spec.api_format,
provider_api_format,
source_model: body_json.get("model").and_then(Value::as_str),
family: same_format_provider_family(spec.family),
body_rules,

View File

@@ -10,6 +10,7 @@ pub(crate) fn build_same_format_upstream_url(
parts: &http::request::Parts,
transport: &GatewayProviderTransportSnapshot,
mapped_model: &str,
provider_api_format: &str,
spec: LocalSameFormatProviderSpec,
upstream_is_stream: bool,
kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>,
@@ -17,7 +18,7 @@ pub(crate) fn build_same_format_upstream_url(
build_same_format_provider_upstream_url_impl(
transport,
SameFormatProviderUpstreamUrlParams {
provider_api_format: spec.api_format,
provider_api_format,
mapped_model,
upstream_is_stream,
request_query: parts.uri.query(),

View File

@@ -28,11 +28,12 @@ pub(crate) use aether_ai_formats::api::{
convert_openai_chat_request_to_openai_responses_request,
convert_openai_chat_response_to_claude_chat, convert_openai_chat_response_to_gemini_chat,
convert_openai_chat_response_to_openai_responses,
convert_openai_responses_response_to_openai_chat, convert_standard_chat_response,
convert_standard_cli_response, copy_request_number_field, copy_request_number_field_as,
core_error_background_report_kind, core_error_default_client_api_format,
core_success_background_report_kind, default_model_for_openai_image_operation, encode_done_sse,
encode_json_sse, encode_kiro_sse_events, estimate_kiro_tokens, extract_openai_text_content,
convert_openai_responses_response_to_openai_chat, convert_request,
convert_standard_chat_response, convert_standard_cli_response, copy_request_number_field,
copy_request_number_field_as, core_error_background_report_kind,
core_error_default_client_api_format, core_success_background_report_kind,
default_model_for_openai_image_operation, encode_done_sse, encode_json_sse,
encode_kiro_sse_events, estimate_kiro_tokens, extract_openai_text_content,
find_kiro_real_thinking_end_tag, find_kiro_real_thinking_end_tag_at_buffer_end,
find_kiro_real_thinking_start_tag, force_upstream_streaming_for_provider,
implicit_sync_finalize_report_kind, is_core_error_finalize_kind,
@@ -75,7 +76,7 @@ pub(crate) use aether_ai_formats::api::{
sync_chat_response_conversion_kind, sync_cli_response_conversion_kind,
transform_provider_private_stream_line, value_as_u64, AiControlPlanRequest,
AiSurfaceFinalizeError, AiSurfaceStreamRewriter, CanonicalStreamFrame, ClaudeClientEmitter,
ClaudeProviderState, ExecutionRuntimeAuthContext, FinalizeStreamRewriteMode,
ClaudeProviderState, ExecutionRuntimeAuthContext, FinalizeStreamRewriteMode, FormatContext,
GeminiClientEmitter, GeminiProviderState, KiroToClaudeCliStreamState, LocalCoreSyncErrorKind,
LocalGeminiFilesSpec, LocalOpenAiImageSpec, LocalOpenAiResponsesSpec,
LocalSameFormatProviderFamily, LocalSameFormatProviderSpec, LocalStandardSourceFamily,
@@ -110,9 +111,10 @@ pub(crate) use aether_ai_formats::api::{
KIRO_ENVELOPE_NAME, KIRO_MAX_THINKING_BUFFER, OPENAI_CHAT_STREAM_PLAN_KIND,
OPENAI_CHAT_STREAM_SUCCESS_REPORT_KIND, OPENAI_CHAT_SYNC_ERROR_REPORT_KIND,
OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND,
OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND,
OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND, OPENAI_RERANK_SYNC_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_ERROR_REPORT_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,