mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
Merge pull request #371 from Kayphoon/feature/embedding-model-support
feat: add embedding and rerank support
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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 => {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user