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

refactor: 抽离 AI pipeline 与调度共享能力逻辑

Browse Source
This commit is contained in:
fawney19
2026-04-10 01:46:14 +08:00
parent b901a6ffc7
commit 5014e2f5fd
255 changed files with 15057 additions and 3115 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/aether-usage-runtime/src/lib.rs
View File

@@ -4,6 +4,7 @@ pub mod queue;
pub mod record;
pub mod report;
pub mod report_context;
mod request_metadata;
pub mod runtime;
pub mod settlement;
pub mod standardized_usage;

36
crates/aether-usage-runtime/src/record.rs
View File

@@ -1,6 +1,7 @@
use aether_data_contracts::repository::usage::UpsertUsageRecord;
use aether_data_contracts::DataLayerError;
use crate::request_metadata::sanitize_usage_request_metadata;
use crate::{UsageEvent, UsageEventType};
pub fn build_upsert_usage_record_from_event(
@@ -56,9 +57,9 @@ pub fn build_upsert_usage_record_from_event(
response_body: data.response_body,
client_response_headers: data.client_response_headers,
client_response_body: data.client_response_body,
request_metadata: data.request_metadata,
request_metadata: sanitize_usage_request_metadata(data.request_metadata),
finalized_at_unix_secs: Some(now_unix_secs),
created_at_unix_secs: Some(now_unix_secs),
created_at_unix_ms: Some(now_unix_secs),
updated_at_unix_secs: now_unix_secs,
})
}
@@ -113,4 +114,35 @@ mod tests {
assert_eq!(record.total_tokens, Some(30));
assert_eq!(record.finalized_at_unix_secs, Some(1_700_000_000));
}
#[test]
fn sanitizes_request_metadata_before_building_upsert_record() {
let record = build_upsert_usage_record_from_event(&UsageEvent {
event_type: UsageEventType::Completed,
request_id: "req-2".to_string(),
timestamp_ms: 1_700_000_000_000,
data: UsageEventData {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
request_metadata: Some(serde_json::json!({
"request_id": "req-2",
"provider_id": "provider-1",
"candidate_id": "cand-2",
"key_name": "upstream-primary",
"billing_snapshot": { "status": "complete" }
})),
..UsageEventData::default()
},
})
.expect("record should build");
assert_eq!(
record.request_metadata,
Some(serde_json::json!({
"candidate_id": "cand-2",
"key_name": "upstream-primary",
"billing_snapshot": { "status": "complete" }
}))
);
}
}

8
crates/aether-usage-runtime/src/report_context.rs
View File

@@ -140,9 +140,9 @@ mod tests {
concurrent_requests: None,
extra_data: None,
required_capabilities: None,
created_at_unix_secs: 1,
started_at_unix_secs: None,
finished_at_unix_secs: None,
created_at_unix_ms: 1,
started_at_unix_ms: None,
finished_at_unix_ms: None,
}
}
@@ -177,7 +177,7 @@ mod tests {
next_poll_at_unix_secs: None,
poll_count: 0,
max_poll_count: 360,
created_at_unix_secs: 1,
created_at_unix_ms: 1,
submitted_at_unix_secs: Some(1),
completed_at_unix_secs: None,
updated_at_unix_secs: 1,

229
crates/aether-usage-runtime/src/request_metadata.rs Normal file
View File

@@ -0,0 +1,229 @@
use aether_contracts::ExecutionPlan;
use serde_json::{Map, Value};
pub(crate) fn build_usage_request_metadata_seed(
plan: &ExecutionPlan,
context: Option<&Map<String, Value>>,
) -> Option<Value> {
let mut metadata = context.cloned().unwrap_or_default();
if !has_non_empty_string(&metadata, "candidate_id") {
if let Some(candidate_id) = plan
.candidate_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
metadata.insert(
"candidate_id".to_string(),
Value::String(candidate_id.to_string()),
);
}
}
sanitize_usage_request_metadata(Some(Value::Object(metadata)))
}
pub(crate) fn merge_usage_request_metadata(
base: Option<Value>,
override_value: Option<Value>,
) -> Option<Value> {
let merged = match (base, override_value) {
(Some(Value::Object(mut base)), Some(Value::Object(override_object))) => {
for (key, value) in override_object {
base.insert(key, value);
}
Some(Value::Object(base))
}
(Some(base), None) => Some(base),
(_, Some(override_value)) => Some(override_value),
(None, None) => None,
};
sanitize_usage_request_metadata(merged)
}
pub(crate) fn sanitize_usage_request_metadata(value: Option<Value>) -> Option<Value> {
let Value::Object(object) = value? else {
return None;
};
let mut filtered = Map::new();
copy_non_empty_string(&object, &mut filtered, "candidate_id");
copy_number(&object, &mut filtered, "candidate_index");
copy_non_empty_string(&object, &mut filtered, "key_name");
copy_non_empty_string(&object, &mut filtered, "trace_id");
copy_non_null_value(&object, &mut filtered, "billing_snapshot");
copy_non_null_value(&object, &mut filtered, "dimensions");
copy_non_null_value(&object, &mut filtered, "billing_rule_snapshot");
copy_non_null_value(&object, &mut filtered, "scheduling_audit");
copy_number(&object, &mut filtered, "rate_multiplier");
copy_bool(&object, &mut filtered, "is_free_tier");
(!filtered.is_empty()).then_some(Value::Object(filtered))
}
fn has_non_empty_string(object: &Map<String, Value>, key: &str) -> bool {
object
.get(key)
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty())
}
fn copy_non_empty_string(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
let Some(value) = source
.get(key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return;
};
target.insert(key.to_string(), Value::String(value.to_string()));
}
fn copy_number(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
let Some(value) = source.get(key).filter(|value| value.is_number()) else {
return;
};
target.insert(key.to_string(), value.clone());
}
fn copy_bool(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
let Some(value) = source.get(key).filter(|value| value.is_boolean()) else {
return;
};
target.insert(key.to_string(), value.clone());
}
fn copy_non_null_value(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
let Some(value) = source.get(key).filter(|value| !value.is_null()) else {
return;
};
target.insert(key.to_string(), value.clone());
}
#[cfg(test)]
mod tests {
use aether_contracts::{ExecutionPlan, RequestBody};
use serde_json::json;
use std::collections::BTreeMap;
use super::{
build_usage_request_metadata_seed, merge_usage_request_metadata,
sanitize_usage_request_metadata,
};
fn sample_plan() -> ExecutionPlan {
ExecutionPlan {
request_id: "req-1".to_string(),
candidate_id: Some("cand-1".to_string()),
provider_name: Some("OpenAI".to_string()),
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
method: "POST".to_string(),
url: "https://example.com/v1/chat/completions".to_string(),
headers: BTreeMap::new(),
content_type: None,
content_encoding: None,
body: RequestBody::from_json(json!({"model": "gpt-5"})),
stream: false,
client_api_format: "openai:chat".to_string(),
provider_api_format: "openai:chat".to_string(),
model_name: Some("gpt-5".to_string()),
proxy: None,
tls_profile: None,
timeouts: None,
}
}
#[test]
fn sanitizes_request_metadata_to_allowlist() {
let metadata = sanitize_usage_request_metadata(Some(json!({
"request_id": "req-1",
"provider_id": "provider-1",
"provider_name": "OpenAI",
"model": "gpt-5",
"candidate_id": "cand-1",
"candidate_index": 2,
"key_name": "upstream-primary",
"trace_id": "trace-1",
"billing_snapshot": {"status": "complete"},
"dimensions": {"total_input_context": 10},
"rate_multiplier": 1.25,
"is_free_tier": false,
"original_headers": {"authorization": "Bearer secret"},
"original_request_body": {"messages": []},
"provider_request_headers": {"authorization": "Bearer secret"},
"upstream_url": "https://example.com/v1/chat/completions"
})))
.expect("metadata should remain");
assert_eq!(
metadata,
json!({
"candidate_id": "cand-1",
"candidate_index": 2,
"key_name": "upstream-primary",
"trace_id": "trace-1",
"billing_snapshot": {"status": "complete"},
"dimensions": {"total_input_context": 10},
"rate_multiplier": 1.25,
"is_free_tier": false
})
);
}
#[test]
fn builds_seed_from_context_and_plan_candidate_id() {
let metadata = build_usage_request_metadata_seed(
&sample_plan(),
Some(
json!({
"request_id": "req-1",
"candidate_index": 0,
"key_name": "upstream-primary",
"provider_id": "provider-1",
"billing_snapshot": {"status": "complete"}
})
.as_object()
.expect("object"),
),
)
.expect("metadata should remain");
assert_eq!(
metadata,
json!({
"candidate_id": "cand-1",
"candidate_index": 0,
"key_name": "upstream-primary",
"billing_snapshot": {"status": "complete"}
})
);
}
#[test]
fn merges_and_filters_request_metadata() {
let metadata = merge_usage_request_metadata(
Some(json!({
"candidate_id": "cand-1",
"request_id": "req-1"
})),
Some(json!({
"candidate_index": 0,
"key_name": "upstream-primary",
"provider_name": "OpenAI"
})),
)
.expect("metadata should remain");
assert_eq!(
metadata,
json!({
"candidate_id": "cand-1",
"candidate_index": 0,
"key_name": "upstream-primary"
})
);
}
}

2
crates/aether-usage-runtime/src/worker.rs
View File

@@ -288,7 +288,7 @@ mod tests {
record.status,
record.billing_status,
record
.created_at_unix_secs
.created_at_unix_ms
.unwrap_or(record.updated_at_unix_secs) as i64,
record.updated_at_unix_secs as i64,
record.finalized_at_unix_secs.map(|value| value as i64),

41
crates/aether-usage-runtime/src/write.rs
View File

@@ -6,6 +6,7 @@ use aether_data_contracts::DataLayerError;
use base64::Engine as _;
use serde_json::{json, Map, Value};
use crate::request_metadata::{build_usage_request_metadata_seed, merge_usage_request_metadata};
use crate::{
map_usage_from_response, GatewayStreamReportRequest, GatewaySyncReportRequest, UsageEvent,
UsageEventData, UsageEventType,
@@ -224,7 +225,10 @@ pub fn build_terminal_usage_event_from_outcome(
response_body: outcome.provider_response.clone(),
client_response_headers: outcome.client_response_headers,
client_response_body: outcome.client_response.clone(),
request_metadata: merge_json_value(outcome.request_metadata, outcome.audit_payload),
request_metadata: merge_usage_request_metadata(
outcome.request_metadata,
outcome.audit_payload,
),
..UsageEventData::default()
};
@@ -323,19 +327,7 @@ fn build_terminal_usage_outcome_base(
provider_response: provider_response.clone(),
client_response_headers,
client_response,
request_metadata: Some(Value::Object(Map::from_iter([
(
"request_id".to_string(),
Value::String(plan.request_id.clone()),
),
(
"candidate_id".to_string(),
plan.candidate_id
.clone()
.map(Value::String)
.unwrap_or(Value::Null),
),
]))),
request_metadata: build_usage_request_metadata_seed(plan, context),
audit_payload: report_context.cloned(),
}
}
@@ -440,7 +432,7 @@ fn build_upsert_usage_record(
client_response_body: data.client_response_body,
request_metadata: data.request_metadata,
finalized_at_unix_secs,
created_at_unix_secs: Some(updated_at_unix_secs),
created_at_unix_ms: Some(updated_at_unix_secs),
updated_at_unix_secs,
})
}
@@ -458,23 +450,6 @@ fn build_base_usage_data(plan: &ExecutionPlan, report_context: Option<&Value>) -
.or_else(|| plan.provider_name.clone())
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| "unknown".to_string());
let mut request_metadata = Map::from_iter([
(
"request_id".to_string(),
Value::String(plan.request_id.clone()),
),
(
"candidate_id".to_string(),
plan.candidate_id
.clone()
.map(Value::String)
.unwrap_or(Value::Null),
),
]);
if let Some(key_name) = context_string(context, "key_name") {
request_metadata.insert("key_name".to_string(), Value::String(key_name));
}
UsageEventData {
user_id: context_string(context, "user_id"),
api_key_id: context_string(context, "api_key_id"),
@@ -517,7 +492,7 @@ fn build_base_usage_data(plan: &ExecutionPlan, report_context: Option<&Value>) -
.or_else(|| Some(headers_to_json(&plan.headers))),
provider_request_body: context_value(context, "provider_request_body")
.or_else(|| plan.body.json_body.clone()),
request_metadata: Some(Value::Object(request_metadata)),
request_metadata: build_usage_request_metadata_seed(plan, context),
..UsageEventData::default()
}
}