i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 09:20:22 +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

8
crates/aether-billing/src/default_rule.rs
View File

@@ -33,12 +33,14 @@ impl DefaultBillingRuleGenerator {
let image_output_price_entries = explicit_image_output_price_entries(pricing_config)
.filter(|entries| !entries.is_empty())
.unwrap_or_default();
let image_output_price_default =
explicit_image_output_price_default(pricing_config).unwrap_or(0.0);
let explicit_image_output_price_default = explicit_image_output_price_default(pricing_config);
let image_output_price_default = explicit_image_output_price_default.unwrap_or(0.0);
let has_image_output_pricing =
!image_output_price_entries.is_empty() || explicit_image_output_price_default.is_some();
if tiers.is_empty()
&& pricing.effective_price_per_request().is_none()
&& image_output_price_entries.is_empty()
&& !has_image_output_pricing
{
return None;
}

315
crates/aether-billing/src/service.rs
View File

@@ -30,6 +30,7 @@ impl BillingService {
pricing: &BillingModelPricingSnapshot,
input: &BillingUsageInput,
) -> Result<BillingComputation, ExpressionEvaluationError> {
let image_output_pricing = image_output_pricing_state(pricing);
let Some(rule) =
DefaultBillingRuleGenerator::generate_for_pricing(pricing, &input.task_type)
else {
@@ -43,7 +44,7 @@ impl BillingService {
rule_name: None,
scope: None,
expression: None,
resolved_dimensions: build_dimensions(input),
resolved_dimensions: build_dimensions(input, image_output_pricing),
resolved_variables: BTreeMap::new(),
cost_breakdown: BTreeMap::new(),
total_cost: 0.0,
@@ -62,7 +63,7 @@ impl BillingService {
});
};
let dims = build_dimensions(input);
let dims = build_dimensions(input, image_output_pricing);
let result = self.engine.evaluate(
&rule.expression,
Some(&rule.variables),
@@ -123,7 +124,10 @@ impl Default for BillingService {
}
}
fn build_dimensions(input: &BillingUsageInput) -> BTreeMap<String, Value> {
fn build_dimensions(
input: &BillingUsageInput,
image_output_pricing: ImageOutputPricingState,
) -> BTreeMap<String, Value> {
let normalized_input_tokens = normalize_input_tokens_for_billing(
input.api_format.as_deref(),
input.input_tokens,
@@ -173,10 +177,28 @@ fn build_dimensions(input: &BillingUsageInput) -> BTreeMap<String, Value> {
("image_count".to_string(), json!(input.image_count.max(0))),
(
"image_count_unmetered".to_string(),
json!(if input.output_tokens > 0 {
0
} else {
json!(if image_output_pricing.enabled {
input.image_count.max(0)
} else {
0
}),
),
(
"image_output_pricing_enabled".to_string(),
json!(image_output_pricing.enabled),
),
(
"image_output_matrix_enabled".to_string(),
json!(image_output_pricing.matrix_enabled),
),
(
"image_output_pricing_mode".to_string(),
json!(if image_output_pricing.matrix_enabled {
"matrix"
} else if image_output_pricing.enabled {
"per_image"
} else {
"none"
}),
),
(
@@ -247,6 +269,66 @@ fn build_dimensions(input: &BillingUsageInput) -> BTreeMap<String, Value> {
out
}
#[derive(Debug, Clone, Copy)]
struct ImageOutputPricingState {
enabled: bool,
matrix_enabled: bool,
}
fn image_output_pricing_state(pricing: &BillingModelPricingSnapshot) -> ImageOutputPricingState {
let matrix_enabled = pricing_has_image_output_matrix(pricing);
let default_enabled = pricing_has_image_output_default_price(pricing);
ImageOutputPricingState {
enabled: matrix_enabled || default_enabled,
matrix_enabled,
}
}
fn pricing_has_image_output_matrix(pricing: &BillingModelPricingSnapshot) -> bool {
let Some(config) = pricing.effective_tiered_pricing() else {
return false;
};
[
"image_output_prices",
"image_output_price_per_image",
"image_output_price_matrix",
"image_prices",
]
.iter()
.any(|key| {
config
.get(key)
.is_some_and(image_price_entries_have_matrix_values)
})
}
fn pricing_has_image_output_default_price(pricing: &BillingModelPricingSnapshot) -> bool {
let Some(config) = pricing.effective_tiered_pricing() else {
return false;
};
config
.get("image_output_price_default")
.or_else(|| config.get("image_price_default"))
.or_else(|| {
config
.get("image_output_prices")
.and_then(|value| value.get("default"))
})
.and_then(Value::as_f64)
.is_some()
}
fn image_price_entries_have_matrix_values(value: &Value) -> bool {
match value {
Value::Object(object) => object.iter().any(|(key, value)| {
!key.eq_ignore_ascii_case("default")
&& (value.as_f64().is_some() || image_price_entries_have_matrix_values(value))
}),
Value::Array(items) => items.iter().any(image_price_entries_have_matrix_values),
_ => false,
}
}
fn now_marker() -> String {
SystemTime::now()
.duration_since(UNIX_EPOCH)
@@ -362,6 +444,227 @@ mod tests {
);
}
#[test]
fn image_token_usage_without_image_output_price_bills_tokens_only() {
let pricing = BillingModelPricingSnapshot {
default_price_per_request: None,
default_tiered_pricing: Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 1.0,
"output_price_per_1m": 2.0
}]
})),
..pricing()
};
let result = BillingService::new()
.calculate(
&pricing,
&BillingUsageInput {
task_type: "image".to_string(),
api_format: Some("openai:image".to_string()),
request_count: 1,
input_tokens: 1_000,
output_tokens: 20_000,
cache_creation_tokens: 0,
cache_creation_ephemeral_5m_tokens: 0,
cache_creation_ephemeral_1h_tokens: 0,
cache_read_tokens: 0,
image_count: 1,
image_size: Some("1024x1024".to_string()),
image_quality: Some("medium".to_string()),
image_output_format: Some("png".to_string()),
cache_ttl_minutes: None,
},
)
.expect("billing should calculate");
assert_eq!(
result
.cost_result
.snapshot
.resolved_dimensions
.get("image_output_pricing_mode"),
Some(&json!("none"))
);
assert_eq!(
result
.cost_result
.snapshot
.resolved_dimensions
.get("image_count_unmetered"),
Some(&json!(0))
);
assert_eq!(
result
.cost_result
.snapshot
.cost_breakdown
.get("image_output_cost"),
Some(&0.0)
);
assert_eq!(result.cost_result.cost, 0.041);
}
#[test]
fn image_default_output_price_adds_image_cost_even_with_token_usage() {
let pricing = BillingModelPricingSnapshot {
default_price_per_request: None,
default_tiered_pricing: Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 1.0,
"output_price_per_1m": 2.0
}],
"image_output_price_default": 0.05
})),
..pricing()
};
let result = BillingService::new()
.calculate(
&pricing,
&BillingUsageInput {
task_type: "image".to_string(),
api_format: Some("openai:image".to_string()),
request_count: 1,
input_tokens: 1_000,
output_tokens: 20_000,
cache_creation_tokens: 0,
cache_creation_ephemeral_5m_tokens: 0,
cache_creation_ephemeral_1h_tokens: 0,
cache_read_tokens: 0,
image_count: 1,
image_size: Some("1024x1024".to_string()),
image_quality: Some("medium".to_string()),
image_output_format: Some("png".to_string()),
cache_ttl_minutes: None,
},
)
.expect("billing should calculate");
assert_eq!(
result
.cost_result
.snapshot
.resolved_dimensions
.get("image_output_pricing_mode"),
Some(&json!("per_image"))
);
assert_eq!(
result
.cost_result
.snapshot
.cost_breakdown
.get("image_output_cost"),
Some(&0.05)
);
assert_eq!(result.cost_result.cost, 0.091);
}
#[test]
fn image_default_output_price_generates_rule_without_token_tiers() {
let pricing = BillingModelPricingSnapshot {
default_price_per_request: None,
default_tiered_pricing: Some(json!({
"image_output_price_default": 0.05
})),
..pricing()
};
let result = BillingService::new()
.calculate(
&pricing,
&BillingUsageInput {
task_type: "image".to_string(),
api_format: Some("openai:image".to_string()),
request_count: 1,
input_tokens: 0,
output_tokens: 0,
cache_creation_tokens: 0,
cache_creation_ephemeral_5m_tokens: 0,
cache_creation_ephemeral_1h_tokens: 0,
cache_read_tokens: 0,
image_count: 2,
image_size: Some("1024x1024".to_string()),
image_quality: Some("medium".to_string()),
image_output_format: Some("png".to_string()),
cache_ttl_minutes: None,
},
)
.expect("billing should calculate");
assert_eq!(result.cost_result.status, BillingSnapshotStatus::Complete);
assert_eq!(
result
.cost_result
.snapshot
.cost_breakdown
.get("image_output_cost"),
Some(&0.1)
);
assert_eq!(result.cost_result.cost, 0.1);
}
#[test]
fn image_token_usage_with_matrix_adds_matrix_image_cost() {
let pricing = BillingModelPricingSnapshot {
default_price_per_request: None,
default_tiered_pricing: Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 1.0,
"output_price_per_1m": 2.0
}],
"image_output_price_default": 0.01,
"image_output_prices": {
"1024x1024": { "medium": 0.05 }
}
})),
..pricing()
};
let result = BillingService::new()
.calculate(
&pricing,
&BillingUsageInput {
task_type: "image".to_string(),
api_format: Some("openai:image".to_string()),
request_count: 1,
input_tokens: 1_000,
output_tokens: 20_000,
cache_creation_tokens: 0,
cache_creation_ephemeral_5m_tokens: 0,
cache_creation_ephemeral_1h_tokens: 0,
cache_read_tokens: 0,
image_count: 1,
image_size: Some("1024x1024".to_string()),
image_quality: Some("medium".to_string()),
image_output_format: Some("png".to_string()),
cache_ttl_minutes: None,
},
)
.expect("billing should calculate");
assert_eq!(
result
.cost_result
.snapshot
.resolved_dimensions
.get("image_output_pricing_mode"),
Some(&json!("matrix"))
);
assert_eq!(
result
.cost_result
.snapshot
.cost_breakdown
.get("image_output_cost"),
Some(&0.05)
);
}
#[test]
fn five_minute_cache_ttl_uses_base_cache_prices() {
let pricing = BillingModelPricingSnapshot {

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 {