feat(billing): 引入 model_id 精确计费查找路径,传播模型 ID 至用量事件与 report context

- UsageEventData 新增 model_id / global_model_id 字段,write.rs seed 结构体同步补充
- report_context 新增 model_id / global_model_id / global_model_name,各 payload 构建时从 candidate 传入
- event_enrichment 优先按 model_id 精确查找计费上下文,回退为按名称多轮查找(保留 NoRule 结果降级逻辑)
- BillingReadRepository 新增 find_model_context_by_model_id,memory/sql 分别实现;SQL 查询重构支持按 provider_model_name 和 mappings 匹配并按优先级排序
- pricing.rs 修复:model_tiered_pricing 为空 tiers 时回退到 default_tiered_pricing
- request_metadata 允许字段列表补充 model_id / global_model_id / global_model_name
- admin usage 路由信息输出及脱敏字段列表同步新增三个 model 相关字段
This commit is contained in:
fawney19
2026-04-24 14:40:28 +08:00
parent f3c9835759
commit 780f09c1a2
20 changed files with 665 additions and 27 deletions

View File

@@ -4,10 +4,23 @@ use aether_usage_runtime::{UsageEvent, UsageEventType};
use async_trait::async_trait;
use serde_json::{Map, Value};
use crate::{BillingModelPricingSnapshot, BillingService, BillingUsageInput};
use crate::{
BillingComputation, BillingModelPricingSnapshot, BillingService, BillingSnapshotStatus,
BillingUsageInput,
};
#[async_trait]
pub trait BillingModelContextLookup: Send + Sync {
async fn find_billing_model_context_by_model_id(
&self,
provider_id: &str,
provider_api_key_id: Option<&str>,
model_id: &str,
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
let _ = (provider_id, provider_api_key_id, model_id);
Ok(None)
}
async fn find_billing_model_context(
&self,
provider_id: &str,
@@ -35,23 +48,78 @@ pub async fn enrich_usage_event_with_billing(
else {
return Ok(());
};
let model_name = event.data.model.trim();
if model_name.is_empty() {
if let Some(model_id) = event
.data
.model_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
if let Some(context) = data
.find_billing_model_context_by_model_id(
provider_id,
event.data.provider_api_key_id.as_deref(),
model_id,
)
.await?
{
let pricing = map_pricing_context(context);
let computation = calculate_billing_computation(&pricing, event)?;
apply_billing_computation(event, computation)?;
return Ok(());
}
}
let mut first_no_rule = None;
for lookup_name in billing_model_lookup_names(&event.data) {
let Some(context) = data
.find_billing_model_context(
provider_id,
event.data.provider_api_key_id.as_deref(),
lookup_name,
)
.await?
else {
continue;
};
let pricing = map_pricing_context(context);
let computation = calculate_billing_computation(&pricing, event)?;
if matches!(
computation.cost_result.status,
BillingSnapshotStatus::NoRule
) {
first_no_rule.get_or_insert(computation);
continue;
}
apply_billing_computation(event, computation)?;
return Ok(());
}
let Some(context) = data
.find_billing_model_context(
provider_id,
event.data.provider_api_key_id.as_deref(),
model_name,
)
.await?
else {
return Ok(());
};
if let Some(computation) = first_no_rule {
apply_billing_computation(event, computation)?;
}
Ok(())
}
let pricing = map_pricing_context(context);
fn billing_model_lookup_names(data: &aether_usage_runtime::UsageEventData) -> Vec<&str> {
let mut names = Vec::new();
for value in [data.target_model.as_deref(), Some(data.model.as_str())]
.into_iter()
.flatten()
{
let value = value.trim();
if !value.is_empty() && !names.contains(&value) {
names.push(value);
}
}
names
}
fn calculate_billing_computation(
pricing: &BillingModelPricingSnapshot,
event: &UsageEvent,
) -> Result<BillingComputation, DataLayerError> {
let input = BillingUsageInput {
task_type: event
.data
@@ -85,11 +153,17 @@ pub async fn enrich_usage_event_with_billing(
cache_ttl_minutes: pricing.provider_api_key_cache_ttl_minutes,
};
let computation = BillingService::new()
.calculate(&pricing, &input)
BillingService::new()
.calculate(pricing, &input)
.map_err(|err| {
DataLayerError::UnexpectedValue(format!("billing calculation failed: {err}"))
})?;
})
}
fn apply_billing_computation(
event: &mut UsageEvent,
computation: BillingComputation,
) -> Result<(), DataLayerError> {
event.data.total_cost_usd = Some(computation.cost_result.cost);
event.data.actual_total_cost_usd = Some(computation.actual_total_cost);
merge_billing_snapshot_metadata(
@@ -97,8 +171,7 @@ pub async fn enrich_usage_event_with_billing(
&computation.cost_result.snapshot,
computation.rate_multiplier,
computation.is_free_tier,
)?;
Ok(())
)
}
fn map_pricing_context(context: StoredBillingModelContext) -> BillingModelPricingSnapshot {
@@ -153,11 +226,22 @@ mod tests {
use super::{enrich_usage_event_with_billing, BillingModelContextLookup};
struct TestLookup {
context: Option<StoredBillingModelContext>,
name_context: Option<StoredBillingModelContext>,
model_id_context: Option<StoredBillingModelContext>,
}
#[async_trait]
impl BillingModelContextLookup for TestLookup {
async fn find_billing_model_context_by_model_id(
&self,
_provider_id: &str,
_provider_api_key_id: Option<&str>,
_model_id: &str,
) -> Result<Option<StoredBillingModelContext>, aether_data_contracts::DataLayerError>
{
Ok(self.model_id_context.clone())
}
async fn find_billing_model_context(
&self,
_provider_id: &str,
@@ -165,14 +249,14 @@ mod tests {
_global_model_name: &str,
) -> Result<Option<StoredBillingModelContext>, aether_data_contracts::DataLayerError>
{
Ok(self.context.clone())
Ok(self.name_context.clone())
}
}
#[tokio::test]
async fn enriches_completed_usage_event_with_billing_snapshot() {
let lookup = TestLookup {
context: Some(
name_context: Some(
StoredBillingModelContext::new(
"provider-1".to_string(),
Some("pay_as_you_go".to_string()),
@@ -192,6 +276,7 @@ mod tests {
)
.expect("billing context should build"),
),
model_id_context: None,
};
let mut event = UsageEvent::new(
UsageEventType::Completed,
@@ -229,4 +314,82 @@ mod tests {
Some("complete")
);
}
#[tokio::test]
async fn enriches_by_provider_model_id_before_name_fallback() {
let blank_name_context = StoredBillingModelContext::new(
"provider-1".to_string(),
Some("pay_as_you_go".to_string()),
Some("key-1".to_string()),
None,
Some(60),
"global-model-blank".to_string(),
"claude-sonnet-4-6".to_string(),
None,
None,
None,
Some("model-blank".to_string()),
Some("claude-sonnet-4-6".to_string()),
None,
None,
None,
)
.expect("blank billing context should build");
let priced_model_context = StoredBillingModelContext::new(
"provider-1".to_string(),
Some("pay_as_you_go".to_string()),
Some("key-1".to_string()),
None,
Some(60),
"global-model-priced".to_string(),
"claude-sonnet-4-6".to_string(),
None,
None,
None,
Some("model-priced".to_string()),
Some("claude-sonnet-4-6".to_string()),
None,
None,
Some(
json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]}),
),
)
.expect("priced billing context should build");
let lookup = TestLookup {
name_context: Some(blank_name_context),
model_id_context: Some(priced_model_context),
};
let mut event = UsageEvent::new(
UsageEventType::Completed,
"req-billing-model-id-1",
UsageEventData {
provider_name: "NekoCode".to_string(),
model: "claude-sonnet-4-6".to_string(),
model_id: Some("model-priced".to_string()),
provider_id: Some("provider-1".to_string()),
provider_api_key_id: Some("key-1".to_string()),
request_type: Some("chat".to_string()),
input_tokens: Some(1_000),
output_tokens: Some(500),
status_code: Some(200),
..UsageEventData::default()
},
);
enrich_usage_event_with_billing(&lookup, &mut event)
.await
.expect("billing should succeed");
assert!(event.data.total_cost_usd.unwrap_or_default() > 0.0);
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|value| value.get("billing_snapshot"))
.and_then(|value| value.get("status"))
.and_then(Value::as_str),
Some("complete")
);
}
}

View File

@@ -24,6 +24,7 @@ impl BillingModelPricingSnapshot {
pub fn effective_tiered_pricing(&self) -> Option<&Value> {
self.model_tiered_pricing
.as_ref()
.filter(|value| has_tiered_pricing_tiers(value))
.or(self.default_tiered_pricing.as_ref())
}
@@ -58,6 +59,67 @@ impl BillingModelPricingSnapshot {
}
}
fn has_tiered_pricing_tiers(value: &Value) -> bool {
value
.get("tiers")
.and_then(Value::as_array)
.is_some_and(|tiers| !tiers.is_empty())
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::BillingModelPricingSnapshot;
fn snapshot(
model_tiered_pricing: Option<serde_json::Value>,
default_tiered_pricing: Option<serde_json::Value>,
) -> BillingModelPricingSnapshot {
BillingModelPricingSnapshot {
provider_id: "provider-1".to_string(),
provider_billing_type: None,
provider_api_key_id: None,
provider_api_key_rate_multipliers: None,
provider_api_key_cache_ttl_minutes: None,
global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-5".to_string(),
global_model_config: None,
default_price_per_request: None,
default_tiered_pricing,
model_id: Some("model-1".to_string()),
model_provider_model_name: Some("gpt-5-upstream".to_string()),
model_config: None,
model_price_per_request: None,
model_tiered_pricing,
}
}
#[test]
fn empty_provider_tiered_pricing_inherits_global_default() {
let default_pricing =
json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]});
let pricing = snapshot(Some(json!({})), Some(default_pricing.clone()));
assert_eq!(pricing.effective_tiered_pricing(), Some(&default_pricing));
let pricing = snapshot(Some(json!({"tiers": []})), Some(default_pricing.clone()));
assert_eq!(pricing.effective_tiered_pricing(), Some(&default_pricing));
}
#[test]
fn populated_provider_tiered_pricing_overrides_global_default() {
let provider_pricing =
json!({"tiers":[{"up_to":null,"input_price_per_1m":1.0,"output_price_per_1m":2.0}]});
let default_pricing =
json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]});
let pricing = snapshot(Some(provider_pricing.clone()), Some(default_pricing));
assert_eq!(pricing.effective_tiered_pricing(), Some(&provider_pricing));
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct BillingUsageInput {
pub task_type: String,