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

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

Browse Source 
feat: add embedding and rerank support
This commit is contained in:
fawney19
2026-05-04 12:56:11 +08:00
committed by GitHub
parent 8dd4135f7c 6bdd7792eb
commit 627eb4f335
87 changed files with 5517 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
apps/aether-gateway/src/ai_serving/api.rs
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(

12
apps/aether-gateway/src/ai_serving/planner/common.rs
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;

18
apps/aether-gateway/src/ai_serving/planner/decision/control_plan.rs
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 => {

8
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/payload.rs
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,

54
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs
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,

9
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request/prepare.rs
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,

4
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/request/body.rs
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,

3
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/request/url.rs
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(),

20
apps/aether-gateway/src/ai_serving/pure/mod.rs
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,

13
apps/aether-gateway/src/api/ai/doubao.rs Normal file
View File

@@ -0,0 +1,13 @@
pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> {
match crate::ai_serving::normalize_api_format_alias(api_format).as_str() {
"doubao:embedding" => Some("doubao:embedding"),
_ => None,
}
}
pub(crate) fn local_path(api_format: &str) -> Option<&'static str> {
match crate::ai_serving::normalize_api_format_alias(api_format).as_str() {
"doubao:embedding" => Some("/v1/embeddings"),
_ => None,
}
}

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

@@ -1,6 +1,7 @@
pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> {
match crate::ai_serving::normalize_api_format_alias(api_format).as_str() {
"gemini:generate_content" => Some("gemini:generate_content"),
"gemini:embedding" => Some("gemini:embedding"),
"gemini:video" => Some("gemini:video"),
"gemini:files" => Some("gemini:files"),
_ => None,
@@ -10,6 +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:video" => Some("/v1beta/models/{model}:predictLongRunning"),
"gemini:files" => Some("/v1beta/files"),
_ => None,

15
apps/aether-gateway/src/api/ai/jina.rs Normal file
View File

@@ -0,0 +1,15 @@
pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> {
match crate::ai_serving::normalize_api_format_alias(api_format).as_str() {
"jina:embedding" => Some("jina:embedding"),
"jina:rerank" => Some("jina:rerank"),
_ => None,
}
}
pub(crate) fn local_path(api_format: &str) -> Option<&'static str> {
match crate::ai_serving::normalize_api_format_alias(api_format).as_str() {
"jina:embedding" => Some("/v1/embeddings"),
"jina:rerank" => Some("/v1/rerank"),
_ => None,
}
}

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

@@ -1,5 +1,7 @@
mod claude;
mod doubao;
mod gemini;
mod jina;
mod openai;
mod registry;

4
apps/aether-gateway/src/api/ai/openai.rs
View File

@@ -1,6 +1,8 @@
pub(crate) fn normalized_signature(api_format: &str) -> Option<&'static str> {
match crate::ai_serving::normalize_api_format_alias(api_format).as_str() {
"openai:chat" => Some("openai:chat"),
"openai:embedding" => Some("openai:embedding"),
"openai:rerank" => Some("openai:rerank"),
"openai:responses" => Some("openai:responses"),
"openai:responses:compact" => Some("openai:responses:compact"),
"openai:image" => Some("openai:image"),
@@ -12,6 +14,8 @@ 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() {
"openai" | "openai:chat" => Some("/v1/chat/completions"),
"openai:embedding" => Some("/v1/embeddings"),
"openai:rerank" => Some("/v1/rerank"),
"openai:responses" => Some("/v1/responses"),
"openai:responses:compact" => Some("/v1/responses/compact"),
"openai:image" => Some("/v1/images/generations"),

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

@@ -1,7 +1,7 @@
use axum::routing::{any, post};
use axum::Router;
use super::{claude, gemini, openai};
use super::{claude, doubao, gemini, jina, openai};
use crate::{handlers::proxy::proxy_request, state::AppState};
// Router registration patterns live here so AI public ingress has a single mount registry.
@@ -9,6 +9,8 @@ use crate::{handlers::proxy::proxy_request, state::AppState};
// which describe operational compatibility surfaces rather than the concrete axum mount list.
const AI_POST_ROUTE_PATTERNS: &[&str] = &[
"/v1/chat/completions",
"/v1/embeddings",
"/v1/rerank",
"/v1/messages",
"/v1/messages/count_tokens",
"/v1/responses",
@@ -45,6 +47,8 @@ pub(crate) fn public_api_format_local_path(api_format: &str) -> &'static str {
openai::local_path(&normalized)
.or_else(|| claude::local_path(&normalized))
.or_else(|| gemini::local_path(&normalized))
.or_else(|| jina::local_path(&normalized))
.or_else(|| doubao::local_path(&normalized))
.unwrap_or("/")
}
@@ -53,6 +57,8 @@ pub(crate) fn normalize_admin_endpoint_signature(api_format: &str) -> Option<&'s
openai::normalized_signature(&normalized)
.or_else(|| claude::normalized_signature(&normalized))
.or_else(|| gemini::normalized_signature(&normalized))
.or_else(|| jina::normalized_signature(&normalized))
.or_else(|| doubao::normalized_signature(&normalized))
}
pub(crate) fn admin_endpoint_signature_parts(
@@ -71,3 +77,26 @@ pub(crate) fn admin_default_body_rules_for_signature(
let _ = provider_type;
Some((normalized_api_format, Vec::new()))
}
#[cfg(test)]
mod tests {
use super::{admin_endpoint_signature_parts, public_api_format_local_path};
#[test]
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"),
("jina:embedding", "jina", "embedding", "/v1/embeddings"),
("doubao:embedding", "doubao", "embedding", "/v1/embeddings"),
("openai:rerank", "openai", "rerank", "/v1/rerank"),
("jina:rerank", "jina", "rerank", "/v1/rerank"),
] {
assert_eq!(
admin_endpoint_signature_parts(api_format),
Some((api_format, family, kind))
);
assert_eq!(public_api_format_local_path(api_format), path);
}
}
}

2
apps/aether-gateway/src/constants.rs
View File

@@ -109,6 +109,8 @@ pub(crate) const RUST_FRONTDOOR_OWNED_ROUTE_PATTERNS: &[&str] = &[
"/v1beta/models",
"/v1beta/models/{path...}",
"/v1/chat/completions",
"/v1/embeddings",
"/v1/rerank",
"/v1/images/generations",
"/v1/images/edits",
"/v1/images/variations",

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

@@ -16,6 +16,22 @@ pub(super) fn classify_ai_public_route(
"openai:chat",
true,
))
} else if method == http::Method::POST && normalized_path == "/v1/embeddings" {
Some(classified(
"ai_public",
"openai",
"embedding",
"openai:embedding",
true,
))
} else if method == http::Method::POST && normalized_path == "/v1/rerank" {
Some(classified(
"ai_public",
"openai",
"rerank",
"openai:rerank",
true,
))
} else if method == http::Method::POST
&& matches!(normalized_path, "/v1/responses" | "/v1/responses/compact")
{

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

@@ -20,6 +20,66 @@ fn classifies_claude_count_tokens_as_non_execution_runtime_public_route() {
assert!(!decision.is_execution_runtime_candidate());
}
#[test]
fn classifies_openai_embeddings_as_embedding_not_chat() {
let headers = headers(&[("authorization", "Bearer sk-test")]);
let uri: Uri = "/v1/embeddings".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("openai"));
assert_eq!(decision.route_kind.as_deref(), Some("embedding"));
assert_ne!(decision.route_kind.as_deref(), Some("chat"));
assert_ne!(decision.route_kind.as_deref(), Some("responses"));
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("openai:embedding")
);
assert!(decision.is_execution_runtime_candidate());
}
#[test]
fn classifies_openai_rerank_as_rerank_not_chat() {
let headers = headers(&[("authorization", "Bearer sk-test")]);
let uri: Uri = "/v1/rerank".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("openai"));
assert_eq!(decision.route_kind.as_deref(), Some("rerank"));
assert_ne!(decision.route_kind.as_deref(), Some("chat"));
assert_ne!(decision.route_kind.as_deref(), Some("embedding"));
assert_eq!(
decision.auth_endpoint_signature.as_deref(),
Some("openai:rerank")
);
assert!(decision.is_execution_runtime_candidate());
}
#[test]
fn classifies_openai_chat_and_responses_separately_from_embedding() {
let headers = headers(&[("authorization", "Bearer sk-test")]);
let chat_uri: Uri = "/v1/chat/completions".parse().expect("uri should parse");
let responses_uri: Uri = "/v1/responses".parse().expect("uri should parse");
let chat = classify_control_route(&http::Method::POST, &chat_uri, &headers)
.expect("chat route should classify");
assert_eq!(chat.route_family.as_deref(), Some("openai"));
assert_eq!(chat.route_kind.as_deref(), Some("chat"));
assert_eq!(chat.auth_endpoint_signature.as_deref(), Some("openai:chat"));
assert_ne!(chat.route_kind.as_deref(), Some("embedding"));
let responses = classify_control_route(&http::Method::POST, &responses_uri, &headers)
.expect("responses route should classify");
assert_eq!(responses.route_family.as_deref(), Some("openai"));
assert_eq!(responses.route_kind.as_deref(), Some("responses"));
assert_eq!(
responses.auth_endpoint_signature.as_deref(),
Some("openai:responses")
);
assert_ne!(responses.route_kind.as_deref(), Some("embedding"));
}
#[test]
fn classifies_models_list_as_claude_when_headers_match() {
let headers = headers(&[

2
apps/aether-gateway/src/frontdoor_loop_guard.rs
View File

@@ -37,6 +37,8 @@ pub(crate) fn frontdoor_self_loop_public_ai_path(path: &str) -> bool {
"/v1/messages"
| "/v1/messages/count_tokens"
| "/v1/chat/completions"
| "/v1/embeddings"
| "/v1/rerank"
| "/v1/responses"
| "/v1/responses/compact"
| "/v1beta/files"

46
apps/aether-gateway/src/handlers/admin/model/global/providers.rs
View File

@@ -8,6 +8,38 @@ use aether_data_contracts::repository::global_models::AdminGlobalModelListQuery;
use serde_json::json;
use std::collections::BTreeMap;
const EMBEDDING_API_FORMATS: &[&str] = &[
"openai:embedding",
"jina:embedding",
"gemini:embedding",
"doubao:embedding",
];
fn json_value_contains_string(value: &serde_json::Value, expected: &str) -> bool {
match value {
serde_json::Value::String(value) => value.trim().eq_ignore_ascii_case(expected),
serde_json::Value::Array(values) => values
.iter()
.any(|value| json_value_contains_string(value, expected)),
serde_json::Value::Object(object) => object
.values()
.any(|value| json_value_contains_string(value, expected)),
_ => false,
}
}
fn json_value_contains_embedding_metadata(value: &serde_json::Value) -> bool {
value
.as_object()
.and_then(|object| object.get("embedding"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
|| json_value_contains_string(value, "embedding")
|| EMBEDDING_API_FORMATS
.iter()
.any(|api_format| json_value_contains_string(value, api_format))
}
pub(crate) async fn build_admin_global_model_providers_payload(
state: &AdminAppState<'_>,
global_model_id: &str,
@@ -53,6 +85,7 @@ pub(crate) async fn build_admin_global_model_providers_payload(
"supports_vision": admin_provider_model_effective_capability(&model, "vision"),
"supports_function_calling": admin_provider_model_effective_capability(&model, "function_calling"),
"supports_streaming": admin_provider_model_effective_capability(&model, "streaming"),
"supports_embedding": admin_provider_model_effective_capability(&model, "embedding"),
"is_active": model.is_active,
}))
})
@@ -123,6 +156,14 @@ pub(crate) async fn build_admin_model_catalog_payload(
.and_then(|value| value.get("streaming"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
let mut supports_embedding = global_model
.supported_capabilities
.as_ref()
.is_some_and(json_value_contains_embedding_metadata)
|| global_model
.config
.as_ref()
.is_some_and(json_value_contains_embedding_metadata);
for model in provider_models {
let Some(provider) = provider_ids.get(&model.provider_id) else {
@@ -143,9 +184,12 @@ pub(crate) async fn build_admin_model_catalog_payload(
admin_provider_model_effective_capability(&model, "function_calling");
let model_supports_streaming =
admin_provider_model_effective_capability(&model, "streaming");
let model_supports_embedding =
admin_provider_model_effective_capability(&model, "embedding");
supports_vision |= model_supports_vision;
supports_function_calling |= model_supports_function_calling;
supports_streaming |= model_supports_streaming;
supports_embedding |= model_supports_embedding;
providers.push(json!({
"provider_id": provider.id,
"provider_name": provider.name,
@@ -162,6 +206,7 @@ pub(crate) async fn build_admin_model_catalog_payload(
"supports_vision": model_supports_vision,
"supports_function_calling": model_supports_function_calling,
"supports_streaming": model_supports_streaming,
"supports_embedding": model_supports_embedding,
"is_active": model.is_active,
}));
}
@@ -189,6 +234,7 @@ pub(crate) async fn build_admin_model_catalog_payload(
"supports_vision": supports_vision,
"supports_function_calling": supports_function_calling,
"supports_streaming": supports_streaming,
"supports_embedding": supports_embedding,
}),
}));
}

52
apps/aether-gateway/src/handlers/admin/model/payloads.rs
View File

@@ -1,6 +1,13 @@
use crate::handlers::admin::shared::unix_secs_to_rfc3339;
use aether_data_contracts::repository::global_models::StoredAdminProviderModel;
const EMBEDDING_API_FORMATS: &[&str] = &[
"openai:embedding",
"jina:embedding",
"gemini:embedding",
"doubao:embedding",
];
pub(crate) fn model_tiered_pricing_first_tier_value(
tiered_pricing: Option<&serde_json::Value>,
field_name: &str,
@@ -26,6 +33,50 @@ fn model_effective_capability(
})
}
fn value_contains_string(value: &serde_json::Value, expected: &str) -> bool {
match value {
serde_json::Value::String(value) => value.trim().eq_ignore_ascii_case(expected),
serde_json::Value::Array(values) => values
.iter()
.any(|value| value_contains_string(value, expected)),
serde_json::Value::Object(object) => object
.values()
.any(|value| value_contains_string(value, expected)),
_ => false,
}
}
fn value_has_true_key(value: &serde_json::Value, key: &str) -> bool {
value
.as_object()
.and_then(|object| object.get(key))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
}
fn value_contains_embedding_metadata(value: &serde_json::Value) -> bool {
value_has_true_key(value, "embedding")
|| value_contains_string(value, "embedding")
|| EMBEDDING_API_FORMATS
.iter()
.any(|api_format| value_contains_string(value, api_format))
}
fn model_effective_embedding_capability(model: &StoredAdminProviderModel) -> bool {
model
.config
.as_ref()
.is_some_and(value_contains_embedding_metadata)
|| model
.global_model_supported_capabilities
.as_ref()
.is_some_and(value_contains_embedding_metadata)
|| model
.global_model_config
.as_ref()
.is_some_and(value_contains_embedding_metadata)
}
pub(crate) fn timestamp_or_now(value: Option<u64>, now_unix_secs: u64) -> serde_json::Value {
unix_secs_to_rfc3339(value.unwrap_or(now_unix_secs))
.map(serde_json::Value::String)
@@ -110,6 +161,7 @@ pub(crate) fn admin_provider_model_effective_capability(
model.global_model_config.as_ref(),
"image_generation",
),
"embedding" => model_effective_embedding_capability(model),
_ => false,
}
}

2
apps/aether-gateway/src/handlers/admin/users/shared.rs
View File

@@ -288,6 +288,7 @@ mod tests {
" OPENAI:RESPONSES ".to_string(),
"claude:messages".to_string(),
"gemini:generate_content".to_string(),
"jina:rerank".to_string(),
"openai:responses".to_string(),
]))
.expect("formats should normalize"),
@@ -295,6 +296,7 @@ mod tests {
"openai:responses".to_string(),
"claude:messages".to_string(),
"gemini:generate_content".to_string(),
"jina:rerank".to_string(),
])
);
}

229
apps/aether-gateway/src/handlers/public/ai_public.rs
View File

@@ -33,6 +33,25 @@ const OPENAI_IMAGE_INPUT_FIDELITY_DETAIL: &str = "input_fidelity 仅支持 low
const OPENAI_IMAGE_OUTPUT_COMPRESSION_DETAIL: &str = "output_compression 必须是 0-100 的整数";
const OPENAI_IMAGE_INVALID_JSON_DETAIL: &str = "图片接口 JSON 请求体无效";
const OPENAI_IMAGE_INVALID_MULTIPART_DETAIL: &str = "图片接口 multipart/form-data 请求体无效";
const OPENAI_EMBEDDING_CONTENT_TYPE_DETAIL: &str =
"Embedding request content-type must be application/json";
const OPENAI_EMBEDDING_INVALID_JSON_DETAIL: &str = "Embedding request JSON body is invalid";
const OPENAI_EMBEDDING_MODEL_REQUIRED_DETAIL: &str = "Embedding request model is required";
const OPENAI_EMBEDDING_INPUT_REQUIRED_DETAIL: &str = "Embedding request input is required";
const OPENAI_EMBEDDING_CHAT_PAYLOAD_DETAIL: &str =
"Embedding request must use input, not chat messages";
const OPENAI_EMBEDDING_STREAM_UNSUPPORTED_DETAIL: &str =
"Embedding requests do not support streaming";
const OPENAI_RERANK_CONTENT_TYPE_DETAIL: &str =
"Rerank request content-type must be application/json";
const OPENAI_RERANK_INVALID_JSON_DETAIL: &str = "Rerank request JSON body is invalid";
const OPENAI_RERANK_MODEL_REQUIRED_DETAIL: &str = "Rerank request model is required";
const OPENAI_RERANK_QUERY_REQUIRED_DETAIL: &str = "Rerank request query is required";
const OPENAI_RERANK_DOCUMENTS_REQUIRED_DETAIL: &str = "Rerank request documents are required";
const OPENAI_RERANK_TOP_N_DETAIL: &str = "Rerank request top_n must be a positive integer";
const OPENAI_RERANK_CHAT_PAYLOAD_DETAIL: &str =
"Rerank request must use query/documents, not chat messages";
const OPENAI_RERANK_STREAM_UNSUPPORTED_DETAIL: &str = "Rerank requests do not support streaming";
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum OpenAiImageOperation {
@@ -78,9 +97,15 @@ pub(crate) fn ai_public_local_requires_buffered_body(
.as_ref()
.is_some_and(|decision| {
decision.route_class.as_deref() == Some("ai_public")
&& decision.route_family.as_deref() == Some("claude")
&& decision.route_kind.as_deref() == Some("count_tokens")
&& request_context.request_method == http::Method::POST
&& ((decision.route_family.as_deref() == Some("claude")
&& decision.route_kind.as_deref() == Some("count_tokens"))
|| (decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("embedding")
&& request_context.request_path == "/v1/embeddings")
|| (decision.route_family.as_deref() == Some("openai")
&& decision.route_kind.as_deref() == Some("rerank")
&& request_context.request_path == "/v1/rerank"))
})
}
@@ -124,14 +149,56 @@ fn maybe_build_local_openai_request_validation_response(
return None;
}
let request_body = request_body?;
if decision.route_kind.as_deref() == Some("chat")
&& request_context.request_path == "/v1/chat/completions"
{
return None;
}
if decision.route_kind.as_deref() == Some("embedding")
&& request_context.request_path == "/v1/embeddings"
{
let Some(request_body) = request_body else {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_EMBEDDING_INVALID_JSON_DETAIL,
));
};
if let Err(detail) = validate_openai_embedding_request(
request_context.request_content_type.as_deref(),
request_body,
) {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
detail,
));
}
return None;
}
if decision.route_kind.as_deref() == Some("rerank")
&& request_context.request_path == "/v1/rerank"
{
let Some(request_body) = request_body else {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
OPENAI_RERANK_INVALID_JSON_DETAIL,
));
};
if let Err(detail) = validate_openai_rerank_request(
request_context.request_content_type.as_deref(),
request_body,
) {
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
detail,
));
}
return None;
}
let request_body = request_body?;
if decision.route_kind.as_deref() != Some("image")
|| !matches!(
request_context.request_path.as_str(),
@@ -293,6 +360,160 @@ fn maybe_build_local_openai_request_validation_response(
None
}
fn validate_openai_embedding_request(
content_type: Option<&str>,
request_body: &Bytes,
) -> Result<(), &'static str> {
if !content_type
.unwrap_or_default()
.to_ascii_lowercase()
.contains("application/json")
{
return Err(OPENAI_EMBEDDING_CONTENT_TYPE_DETAIL);
}
if request_body.is_empty() {
return Err(OPENAI_EMBEDDING_INVALID_JSON_DETAIL);
}
let payload = serde_json::from_slice::<Value>(request_body)
.map_err(|_| OPENAI_EMBEDDING_INVALID_JSON_DETAIL)?;
let object = payload
.as_object()
.ok_or(OPENAI_EMBEDDING_INVALID_JSON_DETAIL)?;
if object.contains_key("messages") {
return Err(OPENAI_EMBEDDING_CHAT_PAYLOAD_DETAIL);
}
if object
.get("stream")
.and_then(value_as_bool)
.unwrap_or(false)
{
return Err(OPENAI_EMBEDDING_STREAM_UNSUPPORTED_DETAIL);
}
if object
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.is_none()
{
return Err(OPENAI_EMBEDDING_MODEL_REQUIRED_DETAIL);
}
let Some(input) = object.get("input") else {
return Err(OPENAI_EMBEDDING_INPUT_REQUIRED_DETAIL);
};
if !embedding_input_is_non_empty(input) {
return Err(OPENAI_EMBEDDING_INPUT_REQUIRED_DETAIL);
}
Ok(())
}
fn validate_openai_rerank_request(
content_type: Option<&str>,
request_body: &Bytes,
) -> Result<(), &'static str> {
if !content_type
.unwrap_or_default()
.to_ascii_lowercase()
.contains("application/json")
{
return Err(OPENAI_RERANK_CONTENT_TYPE_DETAIL);
}
if request_body.is_empty() {
return Err(OPENAI_RERANK_INVALID_JSON_DETAIL);
}
let payload = serde_json::from_slice::<Value>(request_body)
.map_err(|_| OPENAI_RERANK_INVALID_JSON_DETAIL)?;
let object = payload
.as_object()
.ok_or(OPENAI_RERANK_INVALID_JSON_DETAIL)?;
if object.contains_key("messages") {
return Err(OPENAI_RERANK_CHAT_PAYLOAD_DETAIL);
}
if object
.get("stream")
.and_then(value_as_bool)
.unwrap_or(false)
{
return Err(OPENAI_RERANK_STREAM_UNSUPPORTED_DETAIL);
}
if object
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.is_none()
{
return Err(OPENAI_RERANK_MODEL_REQUIRED_DETAIL);
}
if object
.get("query")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.is_none()
{
return Err(OPENAI_RERANK_QUERY_REQUIRED_DETAIL);
}
let Some(documents) = object.get("documents").and_then(Value::as_array) else {
return Err(OPENAI_RERANK_DOCUMENTS_REQUIRED_DETAIL);
};
if documents.is_empty() || documents.iter().any(rerank_document_is_empty) {
return Err(OPENAI_RERANK_DOCUMENTS_REQUIRED_DETAIL);
}
if object
.get("top_n")
.or_else(|| object.get("topN"))
.is_some_and(|value| !positive_json_integer(value))
{
return Err(OPENAI_RERANK_TOP_N_DETAIL);
}
Ok(())
}
fn rerank_document_is_empty(value: &Value) -> bool {
match value {
Value::String(text) => text.trim().is_empty(),
Value::Object(object) => object
.get("text")
.and_then(Value::as_str)
.is_some_and(|text| text.trim().is_empty()),
Value::Null => true,
_ => false,
}
}
fn positive_json_integer(value: &Value) -> bool {
value.as_u64().is_some_and(|number| number > 0)
|| value.as_i64().is_some_and(|number| number > 0)
|| value
.as_str()
.and_then(|text| text.trim().parse::<u64>().ok())
.is_some_and(|number| number > 0)
}
fn embedding_input_is_non_empty(value: &Value) -> bool {
match value {
Value::String(text) => !text.trim().is_empty(),
Value::Array(items) if !items.is_empty() => embedding_array_input_is_non_empty(items),
_ => false,
}
}
fn embedding_array_input_is_non_empty(items: &[Value]) -> bool {
items
.iter()
.all(|item| item.as_str().is_some_and(|text| !text.trim().is_empty()))
|| embedding_token_array_is_non_empty(items)
|| items.iter().all(|item| {
item.as_array()
.is_some_and(|items| embedding_token_array_is_non_empty(items))
})
}
fn embedding_token_array_is_non_empty(items: &[Value]) -> bool {
!items.is_empty() && items.iter().all(|item| item.as_u64().is_some())
}
fn image_request_count(value: &Value) -> Option<u64> {
value
.as_u64()

1
apps/aether-gateway/src/handlers/public/catalog_helpers.rs
View File

@@ -210,6 +210,7 @@ fn serialize_public_catalog_model(model: StoredPublicCatalogModel) -> serde_json
"supports_vision": model.supports_vision,
"supports_function_calling": model.supports_function_calling,
"supports_streaming": model.supports_streaming,
"supports_embedding": model.supports_embedding,
"is_active": model.is_active,
})
}

18
apps/aether-gateway/src/handlers/public/support/models/shared.rs
View File

@@ -17,8 +17,14 @@ pub(crate) fn models_api_format(request_context: &GatewayPublicRequestContext) -
"openai:responses" => Some("openai:responses"),
"openai:responses:compact" => Some("openai:responses:compact"),
"openai:image" => Some("openai:image"),
"openai:embedding" => Some("openai:embedding"),
"openai:rerank" => Some("openai:rerank"),
"claude:messages" => Some("claude:messages"),
"gemini:generate_content" => Some("gemini:generate_content"),
"gemini:embedding" => Some("gemini:embedding"),
"jina:embedding" => Some("jina:embedding"),
"jina:rerank" => Some("jina:rerank"),
"doubao:embedding" => Some("doubao:embedding"),
_ => None,
}
}
@@ -32,6 +38,14 @@ const MODELS_CROSS_FORMAT_QUERY_API_FORMATS: &[&str] = &[
"gemini:generate_content",
];
const MODELS_EMBEDDING_QUERY_API_FORMATS: &[&str] = &[
"openai:embedding",
"jina:embedding",
"gemini:embedding",
"doubao:embedding",
];
const MODELS_RERANK_QUERY_API_FORMATS: &[&str] = &["openai:rerank", "jina:rerank"];
pub(super) fn models_query_api_formats(api_format: &str) -> &'static [&'static str] {
match crate::ai_serving::normalize_api_format_alias(api_format).as_str() {
"openai:chat"
@@ -40,6 +54,10 @@ pub(super) fn models_query_api_formats(api_format: &str) -> &'static [&'static s
| "claude:messages"
| "gemini:generate_content" => MODELS_CROSS_FORMAT_QUERY_API_FORMATS,
"openai:image" => &["openai:image"],
"openai:embedding" | "jina:embedding" | "gemini:embedding" | "doubao:embedding" => {
MODELS_EMBEDDING_QUERY_API_FORMATS
}
"openai:rerank" | "jina:rerank" => MODELS_RERANK_QUERY_API_FORMATS,
_ => &[],
}
}

1
apps/aether-gateway/src/handlers/public/support/user_me_catalog.rs
View File

@@ -362,6 +362,7 @@ pub(super) async fn handle_users_me_providers_get(
"supports_vision": model.supports_vision,
"supports_function_calling": model.supports_function_calling,
"supports_streaming": model.supports_streaming,
"supports_embedding": model.supports_embedding,
})
})
.collect::<Vec<_>>(),

1
apps/aether-gateway/src/model_fetch/tests.rs
View File

@@ -161,6 +161,7 @@ fn sample_provider_model(
Some(global_model_name.to_string()),
None,
None,
None,
Some(json!({ "model_mappings": mappings })),
)
.expect("provider model should build")

149
apps/aether-gateway/src/tests/control/admin/models/global.rs
View File

@@ -450,6 +450,128 @@ async fn gateway_handles_admin_model_catalog_locally_with_trusted_admin_principa
upstream_handle.abort();
}
#[tokio::test]
async fn admin_global_models_include_embedding_capability() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let mut global_model = sample_admin_global_model(
"global-embedding-small",
"text-embedding-3-small",
"Text Embedding 3 Small",
);
global_model.supported_capabilities = Some(json!(["embedding"]));
global_model.config = Some(json!({
"api_formats": ["openai:embedding"],
"dimensions": 1536,
"model_type": "embedding"
}));
let mut provider_model = sample_admin_provider_model(
"model-openai-embedding-small",
"provider-openai",
"global-embedding-small",
"text-embedding-3-small",
);
provider_model.provider_model_mappings = Some(json!([{
"name": "text-embedding-3-small",
"priority": 1,
"api_formats": ["openai:embedding"]
}]));
provider_model.config = Some(json!({
"api_formats": ["openai:embedding"],
"dimensions": 1536,
"model_type": "embedding"
}));
provider_model.supports_streaming = Some(false);
provider_model.global_model_name = Some("text-embedding-3-small".to_string());
provider_model.global_model_display_name = Some("Text Embedding 3 Small".to_string());
provider_model.global_model_supported_capabilities = Some(json!(["embedding"]));
provider_model.global_model_config = global_model.config.clone();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
Vec::new(),
Vec::new(),
));
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new())
.with_admin_global_models(vec![global_model])
.with_admin_provider_models(vec![provider_model]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_global_model_repository_for_tests(global_model_repository),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let list_response = client
.get(format!("{gateway_url}/api/admin/models/global?limit=20"))
.header(crate::constants::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")
.send()
.await
.expect("request should succeed");
assert_eq!(list_response.status(), StatusCode::OK);
let list_payload: serde_json::Value =
list_response.json().await.expect("json body should parse");
assert_eq!(
list_payload["models"][0]["supported_capabilities"],
json!(["embedding"])
);
assert_eq!(
list_payload["models"][0]["config"]["api_formats"],
json!(["openai:embedding"])
);
let catalog_response = client
.get(format!("{gateway_url}/api/admin/models/catalog"))
.header(crate::constants::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")
.send()
.await
.expect("request should succeed");
assert_eq!(catalog_response.status(), StatusCode::OK);
let catalog_payload: serde_json::Value = catalog_response
.json()
.await
.expect("json body should parse");
assert_eq!(
catalog_payload["models"][0]["capabilities"]["supports_embedding"],
true
);
assert_eq!(
catalog_payload["models"][0]["providers"][0]["supports_embedding"],
true
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_returns_service_unavailable_for_admin_model_catalog_without_required_readers() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -859,8 +981,8 @@ async fn gateway_creates_admin_global_model_locally_with_trusted_admin_principal
"output_price_per_1m": 24.0
}]
},
"supported_capabilities": ["streaming", "vision"],
"config": {"streaming": true}
"supported_capabilities": ["streaming", "vision", "embedding"],
"config": {"streaming": true, "api_formats": ["openai:embedding"], "model_type": "embedding"}
}))
.send()
.await
@@ -870,6 +992,14 @@ async fn gateway_creates_admin_global_model_locally_with_trusted_admin_principal
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["name"], "gpt-5-pro");
assert_eq!(payload["display_name"], "GPT 5 Pro");
assert_eq!(
payload["supported_capabilities"],
json!(["streaming", "vision", "embedding"])
);
assert_eq!(
payload["config"]["api_formats"],
json!(["openai:embedding"])
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let created = global_model_repository
@@ -878,6 +1008,10 @@ async fn gateway_creates_admin_global_model_locally_with_trusted_admin_principal
.expect("model lookup should succeed")
.expect("model should exist");
assert_eq!(created.display_name, "GPT 5 Pro");
assert_eq!(
created.supported_capabilities,
Some(json!(["streaming", "vision", "embedding"]))
);
gateway_handle.abort();
upstream_handle.abort();
@@ -926,7 +1060,8 @@ async fn gateway_updates_and_deletes_admin_global_model_locally_with_trusted_adm
.json(&json!({
"display_name": "GPT 5 Updated",
"is_active": false,
"config": {"streaming": false}
"supported_capabilities": ["embedding"],
"config": {"streaming": false, "api_formats": ["openai:embedding"], "dimensions": 1536}
}))
.send()
.await
@@ -939,6 +1074,14 @@ async fn gateway_updates_and_deletes_admin_global_model_locally_with_trusted_adm
.expect("json body should parse");
assert_eq!(update_payload["display_name"], "GPT 5 Updated");
assert_eq!(update_payload["is_active"], false);
assert_eq!(
update_payload["supported_capabilities"],
json!(["embedding"])
);
assert_eq!(
update_payload["config"]["api_formats"],
json!(["openai:embedding"])
);
let delete_response = reqwest::Client::new()
.delete(format!(

77
apps/aether-gateway/src/tests/control/admin/models/provider.rs
View File

@@ -233,7 +233,8 @@ async fn gateway_creates_admin_provider_model_locally_with_trusted_admin_princip
"provider_model_name": "gpt-5-upstream",
"global_model_id": "global-gpt-5",
"supports_vision": true,
"config": {"provider_hint": "gpt-5-upstream"}
"provider_model_mappings": [{"name": "text-embedding-3-small", "priority": 1, "api_formats": ["openai:embedding"]}],
"config": {"provider_hint": "gpt-5-upstream", "api_formats": ["openai:embedding"], "model_type": "embedding"}
}))
.send()
.await
@@ -245,6 +246,11 @@ async fn gateway_creates_admin_provider_model_locally_with_trusted_admin_princip
assert_eq!(payload["global_model_id"], "global-gpt-5");
assert_eq!(payload["provider_model_name"], "gpt-5-upstream");
assert_eq!(payload["effective_supports_vision"], true);
assert_eq!(payload["effective_supports_embedding"], true);
assert_eq!(
payload["provider_model_mappings"][0]["api_formats"],
json!(["openai:embedding"])
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let created = global_model_repository
@@ -258,6 +264,13 @@ async fn gateway_creates_admin_provider_model_locally_with_trusted_admin_princip
.expect("models should read");
assert_eq!(created.len(), 1);
assert_eq!(created[0].provider_model_name, "gpt-5-upstream");
assert_eq!(
created[0]
.config
.as_ref()
.and_then(|value| value.get("api_formats")),
Some(&json!(["openai:embedding"]))
);
gateway_handle.abort();
upstream_handle.abort();
@@ -320,8 +333,10 @@ async fn gateway_updates_and_deletes_admin_provider_model_locally_with_trusted_a
.json(&json!({
"provider_model_name": "gpt-5-mini-upstream",
"global_model_id": "global-gpt-5-mini",
"provider_model_mappings": [{"name": "text-embedding-3-small", "priority": 1, "api_formats": ["openai:embedding"]}],
"supports_streaming": false,
"is_available": false
"is_available": false,
"config": {"api_formats": ["openai:embedding"], "model_type": "embedding"}
}))
.send()
.await
@@ -334,6 +349,7 @@ async fn gateway_updates_and_deletes_admin_provider_model_locally_with_trusted_a
assert_eq!(update_payload["provider_model_name"], "gpt-5-mini-upstream");
assert_eq!(update_payload["global_model_id"], "global-gpt-5-mini");
assert_eq!(update_payload["is_available"], false);
assert_eq!(update_payload["effective_supports_embedding"], true);
let delete_response = reqwest::Client::new()
.delete(format!(
@@ -456,27 +472,37 @@ async fn gateway_handles_admin_provider_available_source_models_locally_with_tru
Vec::new(),
Vec::new(),
));
let mut global_model = sample_admin_global_model(
"global-gpt-5",
"text-embedding-3-small",
"Text Embedding 3 Small",
);
global_model.supported_capabilities = Some(json!(["embedding"]));
global_model.config = Some(json!({"api_formats": ["openai:embedding"]}));
let mut primary_model = sample_admin_provider_model(
"model-openai-gpt5",
"provider-openai",
"global-gpt-5",
"text-embedding-3-small",
);
primary_model.global_model_name = Some("text-embedding-3-small".to_string());
primary_model.global_model_display_name = Some("Text Embedding 3 Small".to_string());
primary_model.global_model_supported_capabilities = Some(json!(["embedding"]));
primary_model.global_model_config = Some(json!({"api_formats": ["openai:embedding"]}));
let mut alternate_model = sample_admin_provider_model(
"model-openai-gpt5-b",
"provider-openai",
"global-gpt-5",
"gpt-5-alt",
);
alternate_model.global_model_name = Some("text-embedding-3-small".to_string());
alternate_model.global_model_display_name = Some("Text Embedding 3 Small".to_string());
alternate_model.global_model_supported_capabilities = Some(json!(["embedding"]));
alternate_model.global_model_config = Some(json!({"api_formats": ["openai:embedding"]}));
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new())
.with_admin_global_models(vec![sample_admin_global_model(
"global-gpt-5",
"gpt-5",
"GPT 5",
)])
.with_admin_provider_models(vec![
sample_admin_provider_model(
"model-openai-gpt5",
"provider-openai",
"global-gpt-5",
"gpt-5-upstream",
),
sample_admin_provider_model(
"model-openai-gpt5-b",
"provider-openai",
"global-gpt-5",
"gpt-5-alt",
),
]),
.with_admin_global_models(vec![global_model])
.with_admin_provider_models(vec![primary_model, alternate_model]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
@@ -507,7 +533,14 @@ async fn gateway_handles_admin_provider_available_source_models_locally_with_tru
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["total"], 1);
assert_eq!(payload["models"][0]["global_model_name"], "gpt-5");
assert_eq!(
payload["models"][0]["global_model_name"],
"text-embedding-3-small"
);
assert_eq!(
payload["models"][0]["capabilities"]["supports_embedding"],
true
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();

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

@@ -1078,6 +1078,12 @@ 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");
assert!(formats
.iter()
.any(|item| item["value"] == "openai:embedding"));
assert!(formats.iter().any(|item| item["value"] == "openai:rerank"));
assert!(formats.iter().any(|item| item["value"] == "jina:embedding"));
assert!(formats.iter().any(|item| item["value"] == "jina:rerank"));
assert!(formats.iter().any(|item| item["value"] == "gemini:video"));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);

1
apps/aether-gateway/src/tests/control/helpers.rs
View File

@@ -330,6 +330,7 @@ pub(super) fn sample_admin_provider_model(
"output_price_per_1m": 20.0,
}]
})),
Some(json!(["streaming", "vision"])),
Some(json!({"streaming": true, "vision": false, "billing": {"currency": "USD"}})),
)
.expect("admin provider model should build")

423
apps/aether-gateway/src/tests/control/proxy/embeddings.rs Normal file
View File

@@ -0,0 +1,423 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use aether_contracts::{ExecutionPlan, ExecutionResult, ResponseBody};
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use http::StatusCode;
use serde_json::json;
use super::super::{
any, build_router_with_state, build_state_with_execution_runtime_override, hash_api_key,
sample_currently_usable_auth_snapshot, sample_endpoint, sample_key, sample_provider,
start_server, AppState, GatewayDataState, InMemoryAuthApiKeySnapshotRepository,
InMemoryProviderCatalogReadRepository, Json, Router,
};
use crate::constants::{
CONTROL_ENDPOINT_SIGNATURE_HEADER, CONTROL_EXECUTION_RUNTIME_HEADER,
CONTROL_ROUTE_FAMILY_HEADER, CONTROL_ROUTE_KIND_HEADER, EXECUTION_PATH_HEADER,
EXECUTION_PATH_LOCAL_AUTH_DENIED,
};
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
fn embedding_success_state(execution_runtime_url: String) -> AppState {
let mut snapshot =
sample_currently_usable_auth_snapshot("key-embedding-success", "user-embedding-success");
snapshot.user_allowed_providers = None;
snapshot.api_key_allowed_providers = None;
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!["text-embedding-3-small".to_string()]);
snapshot.api_key_allowed_models = Some(vec!["text-embedding-3-small".to_string()]);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-embedding-success")),
snapshot,
)]));
let candidate_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
embedding_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider(
"provider-embedding",
"OpenAI Embeddings",
1,
)],
vec![sample_endpoint(
"endpoint-embedding",
"provider-embedding",
"openai:embedding",
"https://api.openai.example",
)],
vec![sample_key(
"key-upstream-embedding",
"provider-embedding",
"openai:embedding",
"sk-upstream-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 embedding_execution_runtime() -> Router {
Router::new().route(
"/v1/execute/sync",
any(|Json(plan): Json<ExecutionPlan>| async move {
assert_embedding_execution_plan(&plan);
Json(embedding_execution_result(&plan))
}),
)
}
fn embedding_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-embedding".to_string(),
provider_name: "OpenAI Embeddings".to_string(),
provider_type: "custom".to_string(),
provider_priority: 1,
provider_is_active: true,
endpoint_id: "endpoint-embedding".to_string(),
endpoint_api_format: "openai:embedding".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("embedding".to_string()),
endpoint_is_active: true,
key_id: "key-upstream-embedding".to_string(),
key_name: "default".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:embedding".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 50,
key_global_priority_by_format: None,
model_id: "model-embedding-small".to_string(),
global_model_id: "global-embedding-small".to_string(),
global_model_name: "text-embedding-3-small".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(false),
model_provider_model_name: "upstream-embedding".to_string(),
model_provider_model_mappings: None,
model_supports_streaming: Some(false),
model_is_active: true,
model_is_available: true,
}
}
fn assert_embedding_execution_plan(plan: &ExecutionPlan) {
assert_eq!(plan.client_api_format, "openai:embedding");
assert_eq!(plan.provider_api_format, "openai:embedding");
assert_eq!(plan.method, "POST");
assert_eq!(plan.url, "https://api.openai.example/v1/embeddings");
assert_eq!(plan.model_name.as_deref(), Some("text-embedding-3-small"));
let body = plan.body.json_body.as_ref().expect("json request body");
assert_eq!(body["model"], "upstream-embedding");
assert!(body.get("input").is_some());
}
fn 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!({
"object": "list",
"model": "upstream-embedding",
"data": [
{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}
],
"usage": {"prompt_tokens": 4, "total_tokens": 4}
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
}
}
#[tokio::test]
async fn embeddings_route_accepts_openai_payload() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(embedding_execution_runtime()).await;
let gateway = build_router_with_state(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-embedding-success")
.json(&json!({
"model": "text-embedding-3-small",
"input": ["hello", "world"]
}))
.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("openai")
);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_KIND_HEADER)
.and_then(|value| value.to_str().ok()),
Some("embedding")
);
assert_ne!(
response
.headers()
.get(CONTROL_ROUTE_KIND_HEADER)
.and_then(|value| value.to_str().ok()),
Some("chat")
);
assert_ne!(
response
.headers()
.get(CONTROL_ROUTE_KIND_HEADER)
.and_then(|value| value.to_str().ok()),
Some("responses")
);
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["data"][0]["object"], "embedding");
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_accepts_all_canonical_input_shapes() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(embedding_execution_runtime()).await;
let gateway = build_router_with_state(embedding_success_state(execution_runtime_url));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
for input in [
json!("hello"),
json!(["hello", "world"]),
json!([1, 2, 3]),
json!([[1, 2], [3, 4]]),
] {
let response = client
.post(format!("{gateway_url}/v1/embeddings"))
.header(http::header::AUTHORIZATION, "Bearer sk-embedding-success")
.json(&json!({
"model": "text-embedding-3-small",
"input": input
}))
.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["data"][0]["embedding"], json!([0.1, 0.2, 0.3]));
}
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn embeddings_route_rejects_invalid_local_payloads() {
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let cases = [
("{", "Embedding request JSON body is invalid"),
(
r#"{"input":"hello"}"#,
"Embedding request model is required",
),
(
r#"{"model":"text-embedding-3-small","input":[]}"#,
"Embedding request input is required",
),
(
r#"{"model":"text-embedding-3-small","messages":[]}"#,
"Embedding request must use input, not chat messages",
),
(
r#"{"model":" ","input":"hello"}"#,
"Embedding request model is required",
),
(
r#"{"model":"text-embedding-3-small","input":[[1],[]]}"#,
"Embedding request input is required",
),
(
r#"{"model":"text-embedding-3-small","input":"hello","stream":true}"#,
"Embedding requests do not support streaming",
),
];
for (body, expected_detail) in cases {
let response = client
.post(format!("{gateway_url}/v1/embeddings"))
.header(http::header::CONTENT_TYPE, "application/json")
.body(body)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_KIND_HEADER)
.and_then(|value| value.to_str().ok()),
Some("embedding")
);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(payload["detail"], expected_detail);
}
gateway_handle.abort();
}
#[tokio::test]
async fn embeddings_route_rejects_non_json_content_type() {
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/embeddings"))
.header(http::header::CONTENT_TYPE, "text/plain")
.body(r#"{"model":"text-embedding-3-small","input":"hello"}"#)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(
payload["detail"],
"Embedding request content-type must be application/json"
);
gateway_handle.abort();
}
#[tokio::test]
async fn embeddings_route_rejects_chat_only_model() {
let mut snapshot = sample_currently_usable_auth_snapshot("key-embedding-1", "user-embedding-1");
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!["text-embedding-3-small".to_string()]);
snapshot.api_key_allowed_models = Some(vec!["text-embedding-3-small".to_string()]);
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-embedding-model-guard")),
snapshot,
)]));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_auth_api_key_data_reader_for_tests(repository),
);
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-embedding-model-guard",
)
.json(&json!({
"model": "gpt-5",
"input": "hello"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::FORBIDDEN);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(payload["error"]["message"], "当前密钥不允许访问模型 gpt-5");
gateway_handle.abort();
}
#[tokio::test]
async fn embeddings_route_rejects_chat_only_api_format() {
let mut snapshot = sample_currently_usable_auth_snapshot("key-embedding-2", "user-embedding-2");
snapshot.user_allowed_api_formats = Some(vec!["openai:chat".to_string()]);
snapshot.api_key_allowed_api_formats = Some(vec!["openai:chat".to_string()]);
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-embedding-format-guard")),
snapshot,
)]));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_auth_api_key_data_reader_for_tests(repository),
);
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-embedding-format-guard",
)
.json(&json!({
"model": "text-embedding-3-small",
"input": "hello"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::FORBIDDEN);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(
payload["error"]["message"],
"当前密钥不允许访问 openai:embedding 格式"
);
gateway_handle.abort();
}

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

@@ -1,2 +1,4 @@
mod embeddings;
mod local_denials;
mod rerank;
mod routing;

326
apps/aether-gateway/src/tests/control/proxy/rerank.rs Normal file
View File

@@ -0,0 +1,326 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use aether_contracts::{ExecutionPlan, ExecutionResult, ResponseBody};
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
use http::StatusCode;
use serde_json::json;
use super::super::{
any, build_router_with_state, build_state_with_execution_runtime_override, hash_api_key,
sample_currently_usable_auth_snapshot, sample_endpoint, sample_key, sample_provider,
start_server, AppState, GatewayDataState, InMemoryAuthApiKeySnapshotRepository,
InMemoryProviderCatalogReadRepository, Json, Router,
};
use crate::constants::{
CONTROL_ENDPOINT_SIGNATURE_HEADER, CONTROL_EXECUTION_RUNTIME_HEADER,
CONTROL_ROUTE_FAMILY_HEADER, CONTROL_ROUTE_KIND_HEADER, EXECUTION_PATH_HEADER,
EXECUTION_PATH_LOCAL_AUTH_DENIED,
};
fn rerank_success_state(execution_runtime_url: String) -> AppState {
let mut snapshot =
sample_currently_usable_auth_snapshot("key-rerank-success", "user-rerank-success");
snapshot.user_allowed_providers = None;
snapshot.api_key_allowed_providers = None;
snapshot.user_allowed_api_formats = Some(vec!["openai:rerank".to_string()]);
snapshot.api_key_allowed_api_formats = Some(vec!["openai:rerank".to_string()]);
snapshot.user_allowed_models = Some(vec!["bge-reranker-base".to_string()]);
snapshot.api_key_allowed_models = Some(vec!["bge-reranker-base".to_string()]);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-rerank-success")),
snapshot,
)]));
let candidate_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
rerank_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-rerank", "OpenAI Rerank", 1)],
vec![sample_endpoint(
"endpoint-rerank",
"provider-rerank",
"openai:rerank",
"https://api.openai.example",
)],
vec![sample_key(
"key-upstream-rerank",
"provider-rerank",
"openai:rerank",
"sk-upstream-rerank",
)],
));
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 rerank_execution_runtime() -> Router {
Router::new().route(
"/v1/execute/sync",
any(|Json(plan): Json<ExecutionPlan>| async move {
assert_rerank_execution_plan(&plan);
Json(rerank_execution_result(&plan))
}),
)
}
fn rerank_candidate_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-rerank".to_string(),
provider_name: "OpenAI Rerank".to_string(),
provider_type: "custom".to_string(),
provider_priority: 1,
provider_is_active: true,
endpoint_id: "endpoint-rerank".to_string(),
endpoint_api_format: "openai:rerank".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("rerank".to_string()),
endpoint_is_active: true,
key_id: "key-upstream-rerank".to_string(),
key_name: "default".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:rerank".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 50,
key_global_priority_by_format: None,
model_id: "model-rerank-base".to_string(),
global_model_id: "global-rerank-base".to_string(),
global_model_name: "bge-reranker-base".to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(false),
model_provider_model_name: "upstream-rerank".to_string(),
model_provider_model_mappings: None,
model_supports_streaming: Some(false),
model_is_active: true,
model_is_available: true,
}
}
fn assert_rerank_execution_plan(plan: &ExecutionPlan) {
assert_eq!(plan.client_api_format, "openai:rerank");
assert_eq!(plan.provider_api_format, "openai:rerank");
assert_eq!(plan.method, "POST");
assert_eq!(plan.url, "https://api.openai.example/v1/rerank");
assert_eq!(plan.model_name.as_deref(), Some("bge-reranker-base"));
let body = plan.body.json_body.as_ref().expect("json request body");
assert_eq!(body["model"], "upstream-rerank");
assert_eq!(body["query"], "hello");
assert_eq!(body["documents"], json!(["hello world", "goodbye"]));
assert_eq!(body["top_n"], 1);
}
fn rerank_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": "upstream-rerank",
"results": [
{"index": 0, "relevance_score": 0.98, "document": {"text": "hello world"}}
],
"usage": {"total_tokens": 8}
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
}
}
#[tokio::test]
async fn rerank_route_accepts_openai_payload() {
let (execution_runtime_url, execution_runtime_handle) =
start_server(rerank_execution_runtime()).await;
let gateway = build_router_with_state(rerank_success_state(execution_runtime_url));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/rerank"))
.header(http::header::AUTHORIZATION, "Bearer sk-rerank-success")
.json(&json!({
"model": "bge-reranker-base",
"query": "hello",
"documents": ["hello world", "goodbye"],
"top_n": 1,
"return_documents": true
}))
.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("openai")
);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_KIND_HEADER)
.and_then(|value| value.to_str().ok()),
Some("rerank")
);
assert_eq!(
response
.headers()
.get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
.and_then(|value| value.to_str().ok()),
Some("openai:rerank")
);
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["results"][0]["index"], 0);
assert_eq!(payload["results"][0]["relevance_score"], 0.98);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn rerank_route_rejects_invalid_local_payloads() {
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let cases = [
("{", "Rerank request JSON body is invalid"),
(
r#"{"query":"hello","documents":["doc"]}"#,
"Rerank request model is required",
),
(
r#"{"model":"bge-reranker-base","documents":["doc"]}"#,
"Rerank request query is required",
),
(
r#"{"model":"bge-reranker-base","query":"hello","documents":[]}"#,
"Rerank request documents are required",
),
(
r#"{"model":"bge-reranker-base","query":"hello","messages":[]}"#,
"Rerank request must use query/documents, not chat messages",
),
(
r#"{"model":"bge-reranker-base","query":"hello","documents":["doc"],"top_n":0}"#,
"Rerank request top_n must be a positive integer",
),
(
r#"{"model":"bge-reranker-base","query":"hello","documents":["doc"],"stream":true}"#,
"Rerank requests do not support streaming",
),
];
for (body, expected_detail) in cases {
let response = client
.post(format!("{gateway_url}/v1/rerank"))
.header(http::header::CONTENT_TYPE, "application/json")
.body(body)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
assert_eq!(
response
.headers()
.get(CONTROL_ROUTE_KIND_HEADER)
.and_then(|value| value.to_str().ok()),
Some("rerank")
);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(payload["detail"], expected_detail);
}
gateway_handle.abort();
}
#[tokio::test]
async fn rerank_route_rejects_non_json_content_type() {
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/rerank"))
.header(http::header::CONTENT_TYPE, "text/plain")
.body(r#"{"model":"bge-reranker-base","query":"hello","documents":["doc"]}"#)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(
payload["detail"],
"Rerank request content-type must be application/json"
);
gateway_handle.abort();
}
#[tokio::test]
async fn rerank_route_rejects_chat_only_api_format() {
let mut snapshot = sample_currently_usable_auth_snapshot("key-rerank-2", "user-rerank-2");
snapshot.user_allowed_api_formats = Some(vec!["openai:chat".to_string()]);
snapshot.api_key_allowed_api_formats = Some(vec!["openai:chat".to_string()]);
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-rerank-format-guard")),
snapshot,
)]));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_auth_api_key_data_reader_for_tests(repository),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/rerank"))
.header(http::header::AUTHORIZATION, "Bearer sk-rerank-format-guard")
.json(&json!({
"model": "bge-reranker-base",
"query": "hello",
"documents": ["doc"]
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::FORBIDDEN);
assert_eq!(
response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok()),
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
);
let payload: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(
payload["error"]["message"],
"当前密钥不允许访问 openai:rerank 格式"
);
gateway_handle.abort();
}

1
apps/aether-gateway/src/tests/frontdoor.rs
View File

@@ -340,6 +340,7 @@ fn sample_public_catalog_model(
Some(true),
Some(true),
Some(true),
Some(false),
true,
)
.expect("public catalog model should build")

88
apps/aether-gateway/src/tests/frontdoor/public_support.rs
View File

@@ -843,6 +843,94 @@ async fn gateway_handles_public_catalog_models_without_proxying_upstream() {
upstream_handle.abort();
}
#[tokio::test]
async fn public_catalog_excludes_unsupported_embedding_provider() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/{*path}",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("proxied"))
}
}),
);
let mut active_embedding = sample_public_catalog_model(
"model-openai-embedding-small",
"provider-openai",
"openai",
"text-embedding-3-small",
"text-embedding-3-small",
"Text Embedding 3 Small",
);
active_embedding.supports_embedding = Some(true);
active_embedding.supports_streaming = Some(false);
let mut inactive_embedding = sample_public_catalog_model(
"model-openai-embedding-inactive",
"provider-openai",
"openai",
"text-embedding-3-large",
"text-embedding-3-large",
"Text Embedding 3 Large",
);
inactive_embedding.supports_embedding = Some(true);
inactive_embedding.is_active = false;
let mut unsupported_embedding = sample_public_catalog_model(
"model-unsupported-embedding",
"provider-openai",
"openai",
"legacy-embedding",
"legacy-embedding",
"Legacy Embedding",
);
unsupported_embedding.supports_embedding = Some(false);
unsupported_embedding.is_active = false;
let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::<StoredPublicGlobalModel>::new())
.with_public_catalog_models(vec![
active_embedding,
inactive_embedding,
unsupported_embedding,
]),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_global_model_reader_for_tests(
global_model_repository,
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/public/models?provider_id=provider-openai&limit=10"
))
.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");
let models = payload.as_array().expect("models should be an array");
assert_eq!(models.len(), 1);
assert_eq!(models[0]["id"], "model-openai-embedding-small");
assert_eq!(models[0]["supports_embedding"], true);
assert_eq!(models[0]["supports_streaming"], false);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_public_catalog_search_models_without_proxying_upstream() {
let upstream_hits = Arc::new(Mutex::new(0usize));