mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
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:
@@ -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")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user