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

feat: add embedding and rerank support

Browse Source
This commit is contained in:
Kayphoon
2026-05-03 17:32:41 +08:00
parent 3e2eca4fd0
commit 5abe664d65
87 changed files with 5520 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all 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<_>>(),