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

Persist ranking metadata through request traces

Browse Source
This commit is contained in:
fawney19
2026-04-27 14:05:19 +08:00
parent f52220a8ec
commit 91db4eefd0
19 changed files with 215 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

61
apps/aether-gateway/src/ai_pipeline/planner/candidate_materialization.rs
View File

@@ -4,6 +4,7 @@ use serde_json::Value;
use uuid::Uuid;
use crate::ai_pipeline::planner::candidate_affinity_cache::remember_scheduler_affinity_for_candidate;
use crate::ai_pipeline::planner::candidate_metadata::append_ranking_metadata_to_object;
use crate::ai_pipeline::planner::candidate_resolution::{
EligibleLocalExecutionCandidate, SkippedLocalExecutionCandidate,
};
@@ -177,34 +178,7 @@ fn local_candidate_extra_data_with_ranking(
}
None => serde_json::Map::new(),
};
object.insert(
"ranking_mode".to_string(),
Value::String(format!("{:?}", ranking.ranking_mode)),
);
object.insert(
"priority_mode".to_string(),
Value::String(format!("{:?}", ranking.priority_mode)),
);
object.insert(
"ranking_index".to_string(),
Value::Number(serde_json::Number::from(ranking.ranking_index as u64)),
);
object.insert(
"priority_slot".to_string(),
Value::Number(serde_json::Number::from(i64::from(ranking.priority_slot))),
);
if let Some(promoted_by) = ranking.promoted_by {
object.insert(
"promoted_by".to_string(),
Value::String(promoted_by.to_string()),
);
}
if let Some(demoted_by) = ranking.demoted_by {
object.insert(
"demoted_by".to_string(),
Value::String(demoted_by.to_string()),
);
}
append_ranking_metadata_to_object(&mut object, ranking);
Some(Value::Object(object))
}
@@ -372,7 +346,10 @@ pub(crate) async fn persist_skipped_local_execution_candidates(
&generated_candidate_id,
required_capabilities,
skipped_candidate.skip_reason,
skipped_candidate.extra_data,
local_candidate_extra_data_with_ranking(
skipped_candidate.extra_data,
skipped_candidate.ranking.as_ref(),
),
error_context,
record_runtime_miss_diagnostic,
)
@@ -641,13 +618,23 @@ mod tests {
"pool-skipped",
Some(json!({ "pool_advanced": {} })),
)),
ranking: None,
extra_data: None,
},
SkippedLocalExecutionCandidate {
candidate: sample_candidate("normal-skipped"),
skip_reason: "key_inactive",
transport: None,
extra_data: None,
ranking: Some(SchedulerRankingOutcome {
original_index: 2,
ranking_index: 1,
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
priority_slot: 9,
promoted_by: None,
demoted_by: Some("cross_format"),
}),
extra_data: Some(json!({ "existing": "value" })),
},
],
"persist skipped should not fail",
@@ -662,5 +649,19 @@ mod tests {
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].key_id.as_deref(), Some("normal-skipped"));
assert_eq!(stored[0].candidate_index, 0);
let extra_data = stored[0]
.extra_data
.as_ref()
.and_then(serde_json::Value::as_object)
.expect("skipped ranking metadata should persist");
assert_eq!(extra_data.get("existing"), Some(&json!("value")));
assert_eq!(
extra_data.get("ranking_mode"),
Some(&json!("CacheAffinity"))
);
assert_eq!(extra_data.get("priority_mode"), Some(&json!("Provider")));
assert_eq!(extra_data.get("ranking_index"), Some(&json!(1)));
assert_eq!(extra_data.get("priority_slot"), Some(&json!(9)));
assert_eq!(extra_data.get("demoted_by"), Some(&json!("cross_format")));
}
}

36
apps/aether-gateway/src/ai_pipeline/planner/candidate_metadata.rs
View File

@@ -1,5 +1,5 @@
use aether_contracts::ProxySnapshot;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use aether_scheduler_core::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRankingOutcome};
use serde_json::{json, Map, Value};
use crate::ai_pipeline::planner::candidate_resolution::EligibleLocalExecutionCandidate;
@@ -26,6 +26,40 @@ pub(crate) struct LocalExecutionCandidateMetadataParts<'a> {
pub(crate) extra_fields: Map<String, Value>,
}
pub(crate) fn append_ranking_metadata_to_object(
object: &mut Map<String, Value>,
ranking: &SchedulerRankingOutcome,
) {
object.insert(
"ranking_mode".to_string(),
Value::String(format!("{:?}", ranking.ranking_mode)),
);
object.insert(
"priority_mode".to_string(),
Value::String(format!("{:?}", ranking.priority_mode)),
);
object.insert(
"ranking_index".to_string(),
Value::Number(serde_json::Number::from(ranking.ranking_index as u64)),
);
object.insert(
"priority_slot".to_string(),
Value::Number(serde_json::Number::from(i64::from(ranking.priority_slot))),
);
if let Some(promoted_by) = ranking.promoted_by {
object.insert(
"promoted_by".to_string(),
Value::String(promoted_by.to_string()),
);
}
if let Some(demoted_by) = ranking.demoted_by {
object.insert(
"demoted_by".to_string(),
Value::String(demoted_by.to_string()),
);
}
}
pub(crate) fn build_request_trace_proxy_value(
transport: Option<&GatewayProviderTransportSnapshot>,
resolved_proxy: Option<&ProxySnapshot>,

3
apps/aether-gateway/src/ai_pipeline/planner/candidate_resolution.rs
View File

@@ -27,6 +27,7 @@ pub(crate) struct SkippedLocalExecutionCandidate {
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(crate) skip_reason: &'static str,
pub(crate) transport: Option<Arc<GatewayProviderTransportSnapshot>>,
pub(crate) ranking: Option<SchedulerRankingOutcome>,
pub(crate) extra_data: Option<serde_json::Value>,
}
@@ -131,6 +132,7 @@ where
candidate,
skip_reason: "transport_snapshot_missing",
transport: None,
ranking: None,
extra_data: None,
});
continue;
@@ -145,6 +147,7 @@ where
candidate,
skip_reason,
transport: Some(transport),
ranking: None,
extra_data: None,
}),
None => selectable.push(EligibleLocalExecutionCandidate {

1
apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/family/candidates.rs
View File

@@ -123,6 +123,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
candidate: item.candidate,
skip_reason: item.skip_reason,
transport: None,
ranking: None,
extra_data: None,
})
.chain(skipped_candidates)

1
apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/family/payload.rs
View File

@@ -102,6 +102,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: resolved.transport.endpoint.header_rules.as_ref(),
body_rules: resolved.transport.endpoint.body_rules.as_ref(),

4
apps/aether-gateway/src/ai_pipeline/planner/pool_scheduler.rs
View File

@@ -449,6 +449,7 @@ fn schedule_pool_group(
candidate,
skip_reason: POOL_ACCOUNT_BLOCKED_SKIP_REASON,
transport: Some(transport),
ranking,
extra_data: None,
});
continue;
@@ -459,6 +460,7 @@ fn schedule_pool_group(
candidate,
skip_reason: POOL_ACCOUNT_EXHAUSTED_SKIP_REASON,
transport: Some(transport),
ranking,
extra_data: None,
});
continue;
@@ -469,6 +471,7 @@ fn schedule_pool_group(
candidate,
skip_reason: POOL_COOLDOWN_SKIP_REASON,
transport: Some(transport),
ranking,
extra_data: None,
});
continue;
@@ -482,6 +485,7 @@ fn schedule_pool_group(
candidate,
skip_reason: POOL_COST_LIMIT_REACHED_SKIP_REASON,
transport: Some(transport),
ranking,
extra_data: None,
});
continue;

6
apps/aether-gateway/src/ai_pipeline/planner/report_context.rs
View File

@@ -1,8 +1,10 @@
use std::collections::BTreeMap;
use aether_scheduler_core::SchedulerRankingOutcome;
use serde_json::{Map, Value};
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_metadata::append_ranking_metadata_to_object;
use crate::orchestration::ExecutionAttemptIdentity;
pub(crate) struct LocalExecutionReportContextParts<'a> {
@@ -23,6 +25,7 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) client_api_format: &'a str,
pub(crate) mapped_model: Option<&'a str>,
pub(crate) candidate_group_id: Option<&'a str>,
pub(crate) ranking: Option<&'a SchedulerRankingOutcome>,
pub(crate) upstream_url: Option<&'a str>,
pub(crate) header_rules: Option<&'a Value>,
pub(crate) body_rules: Option<&'a Value>,
@@ -172,6 +175,9 @@ pub(crate) fn build_local_execution_report_context(
Value::String(candidate_group_id.to_string()),
);
}
if let Some(ranking) = parts.ranking {
append_ranking_metadata_to_object(&mut object, ranking);
}
if let Some(upstream_url) = parts.upstream_url {
object.insert(
"upstream_url".to_string(),

1
apps/aether-gateway/src/ai_pipeline/planner/specialized/files/decision.rs
View File

@@ -83,6 +83,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
mapped_model: None,
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
ranking: eligible.ranking.as_ref(),
upstream_url: None,
header_rules: transport.endpoint.header_rules.as_ref(),
body_rules: transport.endpoint.body_rules.as_ref(),

1
apps/aether-gateway/src/ai_pipeline/planner/specialized/image/decision.rs
View File

@@ -80,6 +80,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: transport.endpoint.header_rules.as_ref(),
body_rules: transport.endpoint.body_rules.as_ref(),

1
apps/aether-gateway/src/ai_pipeline/planner/specialized/image/support.rs
View File

@@ -128,6 +128,7 @@ pub(super) async fn list_local_openai_image_candidate_attempts(
candidate: item.candidate,
skip_reason: item.skip_reason,
transport: None,
ranking: None,
extra_data: None,
})
.collect(),

1
apps/aether-gateway/src/ai_pipeline/planner/specialized/video/decision.rs
View File

@@ -66,6 +66,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
ranking: eligible.ranking.as_ref(),
upstream_url: None,
header_rules: transport.endpoint.header_rules.as_ref(),
body_rules: transport.endpoint.body_rules.as_ref(),

1
apps/aether-gateway/src/ai_pipeline/planner/specialized/video/support.rs
View File

@@ -140,6 +140,7 @@ pub(super) async fn list_local_video_create_candidate_attempts(
candidate: item.candidate,
skip_reason: item.skip_reason,
transport: None,
ranking: None,
extra_data: None,
})
.collect(),

1
apps/aether-gateway/src/ai_pipeline/planner/standard/family/candidates.rs
View File

@@ -154,6 +154,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
candidate: skipped_candidate.candidate,
skip_reason: skipped_candidate.skip_reason,
transport: None,
ranking: None,
extra_data: None,
});
}

1
apps/aether-gateway/src/ai_pipeline/planner/standard/family/payload.rs
View File

@@ -82,6 +82,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: resolved.transport.endpoint.header_rules.as_ref(),
body_rules: resolved.transport.endpoint.body_rules.as_ref(),

1
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/decision/payload.rs
View File

@@ -100,6 +100,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
client_api_format: "openai:chat",
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: resolved.transport.endpoint.header_rules.as_ref(),
body_rules: resolved.transport.endpoint.body_rules.as_ref(),

1
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/plans/candidates.rs
View File

@@ -78,6 +78,7 @@ pub(crate) async fn list_local_openai_chat_candidates(
candidate: skipped_candidate.candidate,
skip_reason: skipped_candidate.skip_reason,
transport: None,
ranking: None,
extra_data: None,
});
}

1
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/responses/decision/payload.rs
View File

@@ -98,6 +98,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: eligible.orchestration.candidate_group_id.as_deref(),
ranking: eligible.ranking.as_ref(),
upstream_url: Some(&resolved.upstream_url),
header_rules: resolved.transport.endpoint.header_rules.as_ref(),
body_rules: resolved.transport.endpoint.body_rules.as_ref(),

1
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/responses/decision/support.rs
View File

@@ -210,6 +210,7 @@ pub(crate) async fn materialize_local_openai_responses_candidate_attempts(
candidate: skipped_candidate.candidate,
skip_reason: skipped_candidate.skip_reason,
transport: None,
ranking: None,
extra_data: None,
});
}