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

feat(billing): add image output pricing and usage tracking

Browse Source
This commit is contained in:
ZheFox
2026-05-18 02:49:56 +08:00
parent 691ccaaa04
commit 0b2a8fafce
9 changed files with 884 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

95
crates/aether-usage-runtime/src/usage_mapper.rs
View File

@@ -33,10 +33,15 @@ impl UsageMapper {
pub fn map_from_response(response: &serde_json::Value, api_format: &str) -> StandardizedUsage {
let family = api_family(api_format);
let Some(usage_value) = resolve_usage_value(response, family.as_str()) else {
return StandardizedUsage::new();
let mut usage = if let Some(usage_value) = resolve_usage_value(response, family.as_str()) {
Self::map(usage_value, api_format, None)
} else {
StandardizedUsage::new()
};
Self::map(usage_value, api_format, None)
if is_openai_image_api(api_format) {
apply_openai_image_response_dimensions(response, &mut usage);
}
usage
}
}
@@ -60,6 +65,53 @@ fn api_family(api_format: &str) -> String {
.to_ascii_lowercase()
}
fn api_kind(api_format: &str) -> String {
api_format
.split(':')
.nth(1)
.unwrap_or_default()
.trim()
.to_ascii_lowercase()
}
fn is_openai_image_api(api_format: &str) -> bool {
api_family(api_format).as_str() == "openai" && api_kind(api_format).as_str() == "image"
}
fn apply_openai_image_response_dimensions(
response: &serde_json::Value,
usage: &mut StandardizedUsage,
) {
let image_count = openai_image_response_image_count(response);
if image_count <= 0 {
return;
}
usage.request_count = image_count;
usage
.dimensions
.insert("image_count".to_string(), serde_json::json!(image_count));
}
fn openai_image_response_image_count(response: &serde_json::Value) -> i64 {
response
.get("data")
.and_then(serde_json::Value::as_array)
.map(|items| items.len() as i64)
.filter(|value| *value > 0)
.or_else(|| image_result_count(response.get("result")))
.unwrap_or(0)
}
fn image_result_count(value: Option<&serde_json::Value>) -> Option<i64> {
match value? {
serde_json::Value::Array(items) => Some(items.len() as i64).filter(|count| *count > 0),
serde_json::Value::Object(object) if !object.is_empty() => Some(1),
serde_json::Value::String(text) if !text.trim().is_empty() => Some(1),
_ => None,
}
}
fn base_mapping(api_format: &str) -> BTreeMap<String, String> {
let mut mapping = BTreeMap::new();
match api_family(api_format).as_str() {
@@ -620,4 +672,41 @@ mod tests {
assert_eq!(usage.output_tokens, 6);
assert_eq!(usage.cache_read_tokens, 2);
}
#[test]
fn maps_openai_image_response_dimensions_without_usage() {
let usage = map_usage_from_response(
&serde_json::json!({
"created": 1_700_000_000,
"data": [
{ "b64_json": "abc" },
{ "url": "https://example.test/image.png" }
]
}),
"openai:image",
);
assert_eq!(usage.request_count, 2);
assert_eq!(usage.dimensions.get("image_count"), Some(&serde_json::json!(2)));
}
#[test]
fn maps_openai_image_response_dimensions_with_native_usage() {
let usage = map_usage_from_response(
&serde_json::json!({
"usage": {
"input_tokens": 11,
"output_tokens": 22,
"total_tokens": 33
},
"data": [{ "b64_json": "abc" }]
}),
"openai:image",
);
assert_eq!(usage.input_tokens, 11);
assert_eq!(usage.output_tokens, 22);
assert_eq!(usage.request_count, 1);
assert_eq!(usage.dimensions.get("image_count"), Some(&serde_json::json!(1)));
}
}

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

@@ -2298,6 +2298,7 @@ fn apply_completed_image_usage_estimate(data: &mut UsageEventData) {
if !usage_event_data_is_image(data) {
return;
}
apply_completed_image_dimensions(data);
if data
.response_body
.as_ref()
@@ -2327,6 +2328,110 @@ fn apply_completed_image_usage_estimate(data: &mut UsageEventData) {
}
}
fn apply_completed_image_dimensions(data: &mut UsageEventData) {
let image_count = usage_dimension_i64(data.request_metadata.as_ref(), "image_count")
.or_else(|| image_response_count(data.response_body.as_ref()))
.or_else(|| image_request_count(data.provider_request_body.as_ref()))
.or_else(|| image_request_count(data.request_body.as_ref()));
if let Some(image_count) = image_count.filter(|value| *value > 0) {
set_usage_dimension_if_absent(data, "image_count", json!(image_count));
}
for (dimension, request_key) in [
("image_size", "size"),
("image_quality", "quality"),
("image_output_format", "output_format"),
] {
if usage_dimension_string(data.request_metadata.as_ref(), dimension).is_some() {
continue;
}
if let Some(value) = image_request_string(data.provider_request_body.as_ref(), request_key)
.or_else(|| image_request_string(data.request_body.as_ref(), request_key))
{
set_usage_dimension_if_absent(data, dimension, json!(value));
}
}
}
fn usage_dimension_i64(metadata: Option<&Value>, key: &str) -> Option<i64> {
metadata
.and_then(Value::as_object)
.and_then(|object| object.get("dimensions"))
.and_then(Value::as_object)
.and_then(|object| object.get(key))
.and_then(|value| {
value
.as_i64()
.or_else(|| value.as_u64().and_then(|number| i64::try_from(number).ok()))
})
}
fn usage_dimension_string(metadata: Option<&Value>, key: &str) -> Option<String> {
metadata
.and_then(Value::as_object)
.and_then(|object| object.get("dimensions"))
.and_then(Value::as_object)
.and_then(|object| object.get(key))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn set_usage_dimension_if_absent(data: &mut UsageEventData, key: &str, value: Value) {
let mut metadata = match data.request_metadata.take() {
Some(Value::Object(object)) => object,
_ => Map::new(),
};
let mut dimensions = match metadata.remove("dimensions") {
Some(Value::Object(object)) => object,
_ => Map::new(),
};
dimensions.entry(key.to_string()).or_insert(value);
metadata.insert("dimensions".to_string(), Value::Object(dimensions));
data.request_metadata = Some(Value::Object(metadata));
}
fn image_response_count(value: Option<&Value>) -> Option<i64> {
let value = value?;
value
.get("data")
.and_then(Value::as_array)
.map(|items| items.len() as i64)
.filter(|count| *count > 0)
.or_else(|| image_result_count(value.get("result")))
}
fn image_result_count(value: Option<&Value>) -> Option<i64> {
match value? {
Value::Array(items) => Some(items.len() as i64).filter(|count| *count > 0),
Value::Object(object) if !object.is_empty() => Some(1),
Value::String(text) if !text.trim().is_empty() => Some(1),
_ => None,
}
}
fn image_request_count(value: Option<&Value>) -> Option<i64> {
value
.and_then(|value| value.get("n"))
.and_then(|value| {
value
.as_i64()
.or_else(|| value.as_u64().and_then(|number| i64::try_from(number).ok()))
})
.filter(|count| *count > 0)
}
fn image_request_string(value: Option<&Value>, key: &str) -> Option<String> {
value
.and_then(|value| value.get(key))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn usage_event_data_is_image(data: &UsageEventData) -> bool {
data.request_type
.as_deref()
@@ -3791,6 +3896,144 @@ mod tests {
);
}
#[test]
fn sync_completed_openai_image_usage_infers_image_dimensions_from_response() {
let request_body = json!({
"model": "gpt-image-2",
"prompt": "draw a small red cube on a clean desk",
"size": "1024x1024",
"quality": "medium",
"output_format": "png"
});
let plan = ExecutionPlan {
request_id: "req-image-sync-completed-dimensions-1".to_string(),
candidate_id: Some("cand-image-sync-completed-dimensions-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/images/generations".to_string(),
headers: BTreeMap::new(),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(request_body.clone()),
stream: false,
client_api_format: "openai:image".to_string(),
provider_api_format: "openai:image".to_string(),
model_name: Some("gpt-image-2".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let payload = GatewaySyncReportRequest {
trace_id: "trace-image-sync-completed-dimensions-1".to_string(),
report_kind: "openai_image_sync_success".to_string(),
report_context: Some(json!({
"client_api_format": "openai:image",
"provider_api_format": "openai:image",
"provider_request_body": request_body
})),
status_code: 200,
headers: BTreeMap::new(),
body_json: Some(json!({
"created": 1_700_000_000,
"data": [{ "b64_json": "abc" }]
})),
client_body_json: None,
body_base64: None,
telemetry: None,
};
let event =
build_sync_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload)
.expect("usage event should build");
let dimensions = event
.data
.request_metadata
.as_ref()
.and_then(|metadata| metadata.get("dimensions"))
.expect("dimensions should exist");
assert_eq!(event.event_type, UsageEventType::Completed);
assert!(event.data.input_tokens.unwrap_or_default() > 0);
assert_eq!(dimensions.get("image_count"), Some(&json!(1)));
assert_eq!(dimensions.get("image_size"), Some(&json!("1024x1024")));
assert_eq!(dimensions.get("image_quality"), Some(&json!("medium")));
assert_eq!(dimensions.get("image_output_format"), Some(&json!("png")));
}
#[test]
fn sync_completed_openai_image_usage_preserves_native_usage_with_image_count() {
let request_body = json!({
"model": "gpt-image-2",
"prompt": "draw a small red cube on a clean desk",
"size": "1024x1024",
"quality": "medium"
});
let plan = ExecutionPlan {
request_id: "req-image-sync-native-usage-1".to_string(),
candidate_id: Some("cand-image-sync-native-usage-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/images/generations".to_string(),
headers: BTreeMap::new(),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(request_body.clone()),
stream: false,
client_api_format: "openai:image".to_string(),
provider_api_format: "openai:image".to_string(),
model_name: Some("gpt-image-2".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let payload = GatewaySyncReportRequest {
trace_id: "trace-image-sync-native-usage-1".to_string(),
report_kind: "openai_image_sync_success".to_string(),
report_context: Some(json!({
"client_api_format": "openai:image",
"provider_api_format": "openai:image",
"provider_request_body": request_body
})),
status_code: 200,
headers: BTreeMap::new(),
body_json: Some(json!({
"usage": {
"input_tokens": 11,
"output_tokens": 22,
"total_tokens": 33
},
"data": [{ "b64_json": "abc" }]
})),
client_body_json: None,
body_base64: None,
telemetry: None,
};
let event =
build_sync_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload)
.expect("usage event should build");
assert_eq!(event.event_type, UsageEventType::Completed);
assert_eq!(event.data.input_tokens, Some(11));
assert_eq!(event.data.output_tokens, Some(22));
assert_eq!(event.data.total_tokens, Some(33));
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|metadata| metadata.get("dimensions"))
.and_then(|dimensions| dimensions.get("image_count")),
Some(&json!(1))
);
}
#[test]
fn stream_terminal_usage_prefers_more_complete_provider_chunks_usage() {
let plan = ExecutionPlan {