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

Merge remote-tracking branch 'upstream/aether-rust-pioneer' into rust

Browse Source
This commit is contained in:
AAEE86
2026-04-10 15:08:32 +08:00
parent b8ce02b4f7 5014e2f5fd
commit 677b8f5acb
255 changed files with 15056 additions and 3205 deletions
Download Patch File Download Diff File Expand all files Collapse all files

147
apps/aether-gateway/src/data/candidate_selection.rs
View File

@@ -4,9 +4,10 @@ use aether_scheduler_core::{
auth_constraints_allow_api_format, build_minimal_candidate_selection,
collect_global_model_names_for_required_capability, normalize_api_format,
resolve_requested_global_model_name, SchedulerAuthConstraints,
SchedulerMinimalCandidateSelectionCandidate,
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
};
use async_trait::async_trait;
use std::collections::BTreeSet;
use super::auth::GatewayAuthApiKeySnapshot;
@@ -62,6 +63,91 @@ pub(crate) async fn read_minimal_candidate_selection(
requested_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
read_minimal_candidate_selection_with_priority_mode(
state,
api_format,
requested_model_name,
require_streaming,
auth_snapshot,
SchedulerPriorityMode::Provider,
)
.await
}
pub(crate) async fn read_minimal_candidate_selection_with_priority_mode(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
priority_mode: SchedulerPriorityMode,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
read_minimal_candidate_selection_with_priority_mode_and_affinity_key(
state,
api_format,
requested_model_name,
require_streaming,
auth_snapshot,
priority_mode,
auth_snapshot_affinity_key(auth_snapshot),
)
.await
}
pub(crate) async fn read_minimal_candidate_selection_with_priority_mode_and_required_capabilities(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
priority_mode: SchedulerPriorityMode,
required_capabilities: Option<&serde_json::Value>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities(
state,
api_format,
requested_model_name,
require_streaming,
auth_snapshot,
priority_mode,
auth_snapshot_affinity_key(auth_snapshot),
required_capabilities,
)
.await
}
pub(crate) async fn read_minimal_candidate_selection_with_priority_mode_and_affinity_key(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
priority_mode: SchedulerPriorityMode,
affinity_key: Option<&str>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities(
state,
api_format,
requested_model_name,
require_streaming,
auth_snapshot,
priority_mode,
affinity_key,
None,
)
.await
}
pub(crate) async fn read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
priority_mode: SchedulerPriorityMode,
affinity_key: Option<&str>,
required_capabilities: Option<&serde_json::Value>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
let normalized_api_format = normalize_api_format(api_format);
if normalized_api_format.is_empty() {
@@ -81,17 +167,16 @@ pub(crate) async fn read_minimal_candidate_selection(
return Ok(Vec::new());
};
let auth_constraints = auth_snapshot.map(auth_snapshot_constraints);
let affinity_key = auth_snapshot
.map(|snapshot| snapshot.api_key_id.trim())
.filter(|value| !value.is_empty());
build_minimal_candidate_selection(
rows,
&normalized_api_format,
requested_model_name,
resolved_global_model_name.as_str(),
require_streaming,
required_capabilities,
auth_constraints.as_ref(),
affinity_key,
priority_mode,
)
}
@@ -128,6 +213,60 @@ pub(crate) async fn read_global_model_names_for_required_capability(
))
}
pub(crate) async fn read_global_model_names_for_api_format(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
) -> Result<Vec<String>, DataLayerError> {
let normalized_api_format = normalize_api_format(api_format);
if normalized_api_format.is_empty() {
return Ok(Vec::new());
}
if !auth_constraints_allow_api_format(
auth_snapshot.map(auth_snapshot_constraints).as_ref(),
&normalized_api_format,
) {
return Ok(Vec::new());
}
let rows = state
.read_minimal_candidate_selection_rows_for_api_format(&normalized_api_format)
.await?;
let auth_constraints = auth_snapshot.map(auth_snapshot_constraints);
let mut model_names = BTreeSet::new();
for row in rows {
if require_streaming && !row.supports_streaming() {
continue;
}
if !aether_scheduler_core::auth_constraints_allow_provider(
auth_constraints.as_ref(),
&row.provider_id,
&row.provider_name,
) {
continue;
}
if !aether_scheduler_core::auth_constraints_allow_model(
auth_constraints.as_ref(),
&row.global_model_name,
&row.global_model_name,
) {
continue;
}
model_names.insert(row.global_model_name);
}
Ok(model_names.into_iter().collect())
}
fn auth_snapshot_affinity_key(auth_snapshot: Option<&GatewayAuthApiKeySnapshot>) -> Option<&str> {
auth_snapshot
.map(|snapshot| snapshot.api_key_id.trim())
.filter(|value| !value.is_empty())
}
fn auth_snapshot_constraints(snapshot: &GatewayAuthApiKeySnapshot) -> SchedulerAuthConstraints {
SchedulerAuthConstraints {
allowed_providers: snapshot

8
apps/aether-gateway/src/data/candidates.rs
View File

@@ -36,7 +36,7 @@ mod tests {
request_id: &str,
candidate_index: i32,
status: RequestCandidateStatus,
started_at_unix_secs: Option<i64>,
started_at_unix_ms: Option<i64>,
latency_ms: Option<i32>,
status_code: Option<i32>,
) -> StoredRequestCandidate {
@@ -62,9 +62,9 @@ mod tests {
Some(1),
None,
None,
100 + i64::from(candidate_index),
started_at_unix_secs,
started_at_unix_secs.map(|value| value + 1),
(100 + i64::from(candidate_index)) * 1_000,
started_at_unix_ms.map(|v| v * 1_000),
started_at_unix_ms.map(|value| (value + 1) * 1_000),
)
.expect("candidate should build")
}

10
apps/aether-gateway/src/data/decision_trace.rs
View File

@@ -101,9 +101,9 @@ mod tests {
Some(1),
None,
Some(serde_json::json!({"cache_1h": true})),
100,
Some(101),
Some(102),
100_000,
Some(101_000),
Some(102_000),
)
.expect("candidate should build")
}
@@ -174,11 +174,15 @@ mod tests {
provider_name: Some("OpenAI".to_string()),
provider_website: Some("https://openai.com".to_string()),
provider_type: Some("custom".to_string()),
provider_priority: Some(0),
provider_keep_priority_on_conversion: Some(false),
endpoint_api_format: Some("openai:chat".to_string()),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
provider_key_name: Some("prod-key".to_string()),
provider_key_auth_type: Some("api_key".to_string()),
provider_key_internal_priority: Some(50),
provider_key_global_priority_by_format: None,
provider_key_capabilities: Some(serde_json::json!({"cache_1h": true})),
provider_key_is_active: Some(true),
}],

34
apps/aether-gateway/src/data/tests.rs
View File

@@ -28,7 +28,9 @@ use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use aether_data_contracts::repository::video_tasks::{
UpsertVideoTask, VideoTaskLookupKey, VideoTaskStatus, VideoTaskWriteRepository,
};
use aether_scheduler_core::{build_minimal_candidate_selection, SchedulerAuthConstraints};
use aether_scheduler_core::{
build_minimal_candidate_selection, SchedulerAuthConstraints, SchedulerPriorityMode,
};
use async_trait::async_trait;
use serde_json::json;
@@ -108,7 +110,7 @@ async fn data_state_find_uses_configured_read_repository() {
next_poll_at_unix_secs: Some(100),
poll_count: 0,
max_poll_count: 360,
created_at_unix_secs: 100,
created_at_unix_ms: 100,
submitted_at_unix_secs: Some(100),
completed_at_unix_secs: None,
updated_at_unix_secs: 100,
@@ -384,9 +386,9 @@ async fn data_state_reads_decision_trace_with_provider_catalog_metadata() {
Some(1),
None,
Some(serde_json::json!({"cache_1h": true})),
100,
Some(101),
Some(102),
100_000,
Some(101_000),
Some(102_000),
)
.expect("candidate should build"),
]));
@@ -472,9 +474,9 @@ async fn data_state_reads_request_audit_bundle_from_multiple_readers() {
Some(1),
None,
Some(serde_json::json!({"cache_1h": true})),
100,
Some(101),
Some(102),
100_000,
Some(101_000),
Some(102_000),
)
.expect("candidate should build"),
]));
@@ -668,8 +670,10 @@ async fn data_state_reads_minimal_candidate_selection_with_auth_filters() {
"gpt-4.1",
"gpt-4.1",
false,
None,
Some(&auth_constraints),
Some(auth_snapshot.api_key_id.as_str()),
SchedulerPriorityMode::Provider,
)
.expect("selection should read");
@@ -722,7 +726,7 @@ async fn maps_openai_video_task_repository_row_into_read_response() {
next_poll_at_unix_secs: Some(120),
poll_count: 1,
max_poll_count: 360,
created_at_unix_secs: 100,
created_at_unix_ms: 100,
submitted_at_unix_secs: Some(100),
completed_at_unix_secs: None,
updated_at_unix_secs: 120,
@@ -781,7 +785,7 @@ async fn maps_gemini_video_task_repository_row_into_read_response() {
next_poll_at_unix_secs: None,
poll_count: 4,
max_poll_count: 360,
created_at_unix_secs: 100,
created_at_unix_ms: 100,
submitted_at_unix_secs: Some(100),
completed_at_unix_secs: Some(120),
updated_at_unix_secs: 120,
@@ -843,7 +847,7 @@ async fn data_state_write_uses_configured_shadow_result_writer() {
match_status: ShadowResultMatchStatus::Pending,
status_code: Some(200),
error_message: None,
created_at_unix_secs: 100,
created_at_unix_ms: 100,
updated_at_unix_secs: 100,
})
.await
@@ -906,7 +910,7 @@ async fn data_state_records_shadow_result_samples_and_merges_match_status() {
.expect("second stored result should exist");
assert_eq!(second.match_status, ShadowResultMatchStatus::Match);
assert_eq!(second.created_at_unix_secs, 100);
assert_eq!(second.created_at_unix_ms, 100);
assert_eq!(second.updated_at_unix_secs, 200);
assert_eq!(second.request_id.as_deref(), Some("req-1"));
}
@@ -986,7 +990,7 @@ fn sample_request_candidate(
request_id: &str,
candidate_index: i32,
status: RequestCandidateStatus,
started_at_unix_secs: Option<i64>,
started_at_unix_ms: Option<i64>,
latency_ms: Option<i32>,
status_code: Option<i32>,
) -> StoredRequestCandidate {
@@ -1013,8 +1017,8 @@ fn sample_request_candidate(
None,
None,
100 + i64::from(candidate_index),
started_at_unix_secs,
started_at_unix_secs.map(|value| value + 1),
started_at_unix_ms,
started_at_unix_ms.map(|value| value + 1),
)
.expect("candidate should build")
}