i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-03 01:40:21 +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

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();
}