mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
feat: 扩展 Rust gateway 全功能模块,新增 billing/crypto/wallet crate 及完整数据层
- 新增 aether-billing、aether-crypto、aether-wallet 独立 crate - aether-data 扩展 repository 层:announcements、auth_modules、billing、 candidate_selection、gemini_file_mappings、global_models、management_tokens、 oauth_providers、proxy_nodes、quota、users、wallet 等模块 - aether-gateway 新增 api/auth/billing/control/middleware/scheduler/usage/ video_tasks/hooks/maintenance/model_fetch/provider_transport 等功能模块 - 重构 executor decision 和 gateway state 为模块目录结构 - 新增 gateway router、frontdoor 路由层及对应测试 - Python 侧 API 路由重构,新增 compat/support 模块 - 前端 Logo 组件更新及 Provider 管理页面调整
This commit is contained in:
12
crates/aether-billing/Cargo.toml
Normal file
12
crates/aether-billing/Cargo.toml
Normal file
@@ -0,0 +1,12 @@
|
||||
[package]
|
||||
name = "aether-billing"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
description = "Shared billing domain core for Aether Rust migration"
|
||||
|
||||
[dependencies]
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
216
crates/aether-billing/src/default_rule.rs
Normal file
216
crates/aether-billing/src/default_rule.rs
Normal file
@@ -0,0 +1,216 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use crate::pricing::BillingModelPricingSnapshot;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct VirtualBillingRule {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub task_type: String,
|
||||
pub expression: String,
|
||||
pub variables: BTreeMap<String, Value>,
|
||||
pub dimension_mappings: BTreeMap<String, Value>,
|
||||
pub scope: String,
|
||||
}
|
||||
|
||||
pub struct DefaultBillingRuleGenerator;
|
||||
|
||||
impl DefaultBillingRuleGenerator {
|
||||
pub fn generate_for_pricing(
|
||||
pricing: &BillingModelPricingSnapshot,
|
||||
task_type: &str,
|
||||
) -> Option<VirtualBillingRule> {
|
||||
let tiers = pricing
|
||||
.effective_tiered_pricing()
|
||||
.and_then(|value| value.get("tiers"))
|
||||
.and_then(Value::as_array)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
|
||||
if tiers.is_empty() && pricing.effective_price_per_request().is_none() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let first_tier = tiers.first().cloned().unwrap_or_else(|| json!({}));
|
||||
let base_input_price = tier_value(&first_tier, "input_price_per_1m", 0.0);
|
||||
let base_output_price = tier_value(&first_tier, "output_price_per_1m", 0.0);
|
||||
let base_cache_creation_price =
|
||||
tier_value_with_fallback(&first_tier, "cache_creation_price_per_1m", 1.25);
|
||||
let base_cache_read_price =
|
||||
tier_value_with_fallback(&first_tier, "cache_read_price_per_1m", 0.1);
|
||||
let base_request_price = pricing.effective_price_per_request().unwrap_or(0.0);
|
||||
|
||||
let mut variables = BTreeMap::new();
|
||||
variables.insert("input_price_per_1m".to_string(), json!(base_input_price));
|
||||
variables.insert("output_price_per_1m".to_string(), json!(base_output_price));
|
||||
variables.insert(
|
||||
"cache_creation_price_per_1m".to_string(),
|
||||
json!(base_cache_creation_price),
|
||||
);
|
||||
variables.insert(
|
||||
"cache_read_price_per_1m".to_string(),
|
||||
json!(base_cache_read_price),
|
||||
);
|
||||
variables.insert("price_per_request".to_string(), json!(base_request_price));
|
||||
|
||||
let mut dimension_mappings = BTreeMap::new();
|
||||
for (name, key, default) in [
|
||||
("input_tokens", "input_tokens", json!(0)),
|
||||
("output_tokens", "output_tokens", json!(0)),
|
||||
("cache_creation_tokens", "cache_creation_tokens", json!(0)),
|
||||
("cache_read_tokens", "cache_read_tokens", json!(0)),
|
||||
("request_count", "request_count", json!(1)),
|
||||
] {
|
||||
dimension_mappings.insert(
|
||||
name.to_string(),
|
||||
json!({
|
||||
"source": "dimension",
|
||||
"key": key,
|
||||
"required": false,
|
||||
"allow_zero": true,
|
||||
"default": default,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
for (name, expression) in [
|
||||
("input_cost", "input_tokens * input_price_per_1m / 1000000"),
|
||||
(
|
||||
"output_cost",
|
||||
"output_tokens * output_price_per_1m / 1000000",
|
||||
),
|
||||
(
|
||||
"cache_creation_cost",
|
||||
"cache_creation_tokens * cache_creation_price_per_1m / 1000000",
|
||||
),
|
||||
(
|
||||
"cache_read_cost",
|
||||
"cache_read_tokens * cache_read_price_per_1m / 1000000",
|
||||
),
|
||||
("request_cost", "request_count * price_per_request"),
|
||||
] {
|
||||
dimension_mappings.insert(
|
||||
name.to_string(),
|
||||
json!({
|
||||
"source": "computed",
|
||||
"expression": expression,
|
||||
"required": false,
|
||||
"default": 0,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
if !tiers.is_empty() {
|
||||
dimension_mappings.insert(
|
||||
"input_price_per_1m".to_string(),
|
||||
json!({
|
||||
"source": "tiered",
|
||||
"tier_key": "total_input_context",
|
||||
"allow_zero": true,
|
||||
"tiers": build_tier_entries(&tiers, "input_price_per_1m", None, false),
|
||||
"default": base_input_price,
|
||||
}),
|
||||
);
|
||||
dimension_mappings.insert(
|
||||
"output_price_per_1m".to_string(),
|
||||
json!({
|
||||
"source": "tiered",
|
||||
"tier_key": "total_input_context",
|
||||
"allow_zero": true,
|
||||
"tiers": build_tier_entries(&tiers, "output_price_per_1m", None, false),
|
||||
"default": base_output_price,
|
||||
}),
|
||||
);
|
||||
dimension_mappings.insert(
|
||||
"cache_creation_price_per_1m".to_string(),
|
||||
json!({
|
||||
"source": "tiered",
|
||||
"tier_key": "total_input_context",
|
||||
"allow_zero": true,
|
||||
"ttl_key": "cache_ttl_minutes",
|
||||
"ttl_value_key": "cache_creation_price_per_1m",
|
||||
"tiers": build_tier_entries(&tiers, "cache_creation_price_per_1m", Some(1.25), true),
|
||||
"default": base_cache_creation_price,
|
||||
}),
|
||||
);
|
||||
dimension_mappings.insert(
|
||||
"cache_read_price_per_1m".to_string(),
|
||||
json!({
|
||||
"source": "tiered",
|
||||
"tier_key": "total_input_context",
|
||||
"allow_zero": true,
|
||||
"ttl_key": "cache_ttl_minutes",
|
||||
"ttl_value_key": "cache_read_price_per_1m",
|
||||
"tiers": build_tier_entries(&tiers, "cache_read_price_per_1m", Some(0.1), true),
|
||||
"default": base_cache_read_price,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
Some(VirtualBillingRule {
|
||||
id: "__default__".to_string(),
|
||||
name: format!("Default rule for {}", pricing.global_model_name),
|
||||
task_type: normalize_task_type(task_type).to_string(),
|
||||
expression:
|
||||
"input_cost + output_cost + cache_creation_cost + cache_read_cost + request_cost"
|
||||
.to_string(),
|
||||
variables,
|
||||
dimension_mappings,
|
||||
scope: "default".to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn normalize_task_type(task_type: &str) -> &str {
|
||||
if task_type.trim().eq_ignore_ascii_case("cli") {
|
||||
"chat"
|
||||
} else {
|
||||
task_type.trim()
|
||||
}
|
||||
}
|
||||
|
||||
fn tier_value(tier: &Value, key: &str, default: f64) -> f64 {
|
||||
tier.get(key).and_then(Value::as_f64).unwrap_or(default)
|
||||
}
|
||||
|
||||
fn tier_value_with_fallback(tier: &Value, key: &str, default_multiplier: f64) -> f64 {
|
||||
if let Some(value) = tier.get(key).and_then(Value::as_f64) {
|
||||
return value;
|
||||
}
|
||||
tier.get("input_price_per_1m")
|
||||
.and_then(Value::as_f64)
|
||||
.map(|value| value * default_multiplier)
|
||||
.unwrap_or(0.0)
|
||||
}
|
||||
|
||||
fn build_tier_entries(
|
||||
tiers: &[Value],
|
||||
key: &str,
|
||||
default_multiplier: Option<f64>,
|
||||
include_cache_ttl_pricing: bool,
|
||||
) -> Vec<Value> {
|
||||
tiers
|
||||
.iter()
|
||||
.map(|tier| {
|
||||
let mut value = serde_json::Map::new();
|
||||
value.insert(
|
||||
"up_to".to_string(),
|
||||
tier.get("up_to").cloned().unwrap_or(Value::Null),
|
||||
);
|
||||
let resolved = match default_multiplier {
|
||||
Some(multiplier) => Value::from(tier_value_with_fallback(tier, key, multiplier)),
|
||||
None => Value::from(tier_value(tier, key, 0.0)),
|
||||
};
|
||||
value.insert("value".to_string(), resolved);
|
||||
if include_cache_ttl_pricing {
|
||||
if let Some(ttl_pricing) = tier.get("cache_ttl_pricing").cloned() {
|
||||
value.insert("cache_ttl_pricing".to_string(), ttl_pricing);
|
||||
}
|
||||
}
|
||||
Value::Object(value)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
860
crates/aether-billing/src/formula_engine.rs
Normal file
860
crates/aether-billing/src/formula_engine.rs
Normal file
@@ -0,0 +1,860 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::precision::quantize_cost;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum FormulaEvaluationStatus {
|
||||
Complete,
|
||||
Incomplete,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct FormulaEvaluationResult {
|
||||
pub status: FormulaEvaluationStatus,
|
||||
pub cost: f64,
|
||||
pub resolved_dimensions: BTreeMap<String, serde_json::Value>,
|
||||
pub resolved_variables: BTreeMap<String, serde_json::Value>,
|
||||
pub cost_breakdown: BTreeMap<String, f64>,
|
||||
pub tier_index: Option<i64>,
|
||||
pub tier_info: Option<serde_json::Value>,
|
||||
pub missing_required: Vec<String>,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum UnsafeExpressionError {
|
||||
#[error("unsupported expression syntax: {0}")]
|
||||
Unsupported(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum ExpressionEvaluationError {
|
||||
#[error("expression evaluation failed: {0}")]
|
||||
Failed(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
#[error("missing required dimensions: {missing_required:?}")]
|
||||
pub struct BillingIncompleteError {
|
||||
pub missing_required: Vec<String>,
|
||||
}
|
||||
|
||||
pub struct FormulaEngine;
|
||||
|
||||
type TierMetadata = Option<(i64, serde_json::Value)>;
|
||||
type MappingResolution = Result<(serde_json::Value, bool, TierMetadata), ExpressionEvaluationError>;
|
||||
|
||||
impl Default for FormulaEngine {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl FormulaEngine {
|
||||
pub fn new() -> Self {
|
||||
Self
|
||||
}
|
||||
|
||||
pub fn evaluate(
|
||||
&self,
|
||||
expression: &str,
|
||||
variables: Option<&BTreeMap<String, serde_json::Value>>,
|
||||
dimensions: Option<&BTreeMap<String, serde_json::Value>>,
|
||||
dimension_mappings: Option<&BTreeMap<String, serde_json::Value>>,
|
||||
strict_mode: bool,
|
||||
) -> Result<FormulaEvaluationResult, ExpressionEvaluationError> {
|
||||
let dims = dimensions.cloned().unwrap_or_default();
|
||||
let mut resolved = variables.cloned().unwrap_or_default();
|
||||
let mut missing_required = Vec::new();
|
||||
let mut tier_index = None;
|
||||
let mut tier_info = None;
|
||||
let mut computed = BTreeMap::new();
|
||||
|
||||
if let Some(mappings) = dimension_mappings {
|
||||
for (var_name, mapping) in mappings {
|
||||
let source = mapping
|
||||
.get("source")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("constant")
|
||||
.to_ascii_lowercase();
|
||||
if source == "computed" {
|
||||
computed.insert(var_name.clone(), mapping.clone());
|
||||
continue;
|
||||
}
|
||||
if source == "constant" && resolved.contains_key(var_name) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let (value, is_missing, tier_meta) = resolve_mapping(var_name, mapping, &dims)?;
|
||||
if let Some((idx, info)) = tier_meta {
|
||||
if tier_index.is_none() {
|
||||
tier_index = Some(idx);
|
||||
tier_info = Some(info);
|
||||
}
|
||||
}
|
||||
if is_missing {
|
||||
missing_required.push(var_name.clone());
|
||||
continue;
|
||||
}
|
||||
resolved.insert(var_name.clone(), value);
|
||||
}
|
||||
}
|
||||
|
||||
if !computed.is_empty() {
|
||||
let mut unresolved = computed;
|
||||
for _ in 0..std::cmp::max(4, unresolved.len() + 1) {
|
||||
let mut progressed = false;
|
||||
let names: Vec<String> = unresolved.keys().cloned().collect();
|
||||
for var_name in names {
|
||||
let mapping = match unresolved.get(&var_name).cloned() {
|
||||
Some(value) => value,
|
||||
None => continue,
|
||||
};
|
||||
let (status, value) = try_resolve_computed(&mapping, &dims, &resolved)?;
|
||||
match status {
|
||||
ComputedStatus::Pending => {}
|
||||
ComputedStatus::MissingRequired => {
|
||||
missing_required.push(var_name.clone());
|
||||
unresolved.remove(&var_name);
|
||||
}
|
||||
ComputedStatus::Ok | ComputedStatus::Defaulted => {
|
||||
resolved.insert(var_name.clone(), value);
|
||||
unresolved.remove(&var_name);
|
||||
progressed = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
if !progressed {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
for (var_name, mapping) in unresolved {
|
||||
let required = mapping
|
||||
.get("required")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
if required {
|
||||
missing_required.push(var_name);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !missing_required.is_empty() {
|
||||
if strict_mode {
|
||||
return Err(ExpressionEvaluationError::Failed(
|
||||
BillingIncompleteError { missing_required }.to_string(),
|
||||
));
|
||||
}
|
||||
return Ok(FormulaEvaluationResult {
|
||||
status: FormulaEvaluationStatus::Incomplete,
|
||||
cost: 0.0,
|
||||
resolved_dimensions: dims,
|
||||
resolved_variables: resolved,
|
||||
cost_breakdown: BTreeMap::new(),
|
||||
tier_index,
|
||||
tier_info,
|
||||
missing_required,
|
||||
error: None,
|
||||
});
|
||||
}
|
||||
|
||||
let cost = evaluate_expression(expression, &resolved)
|
||||
.map_err(|err| ExpressionEvaluationError::Failed(err.to_string()))?;
|
||||
if cost < 0.0 {
|
||||
return Ok(FormulaEvaluationResult {
|
||||
status: FormulaEvaluationStatus::Incomplete,
|
||||
cost: 0.0,
|
||||
resolved_dimensions: dims,
|
||||
resolved_variables: resolved,
|
||||
cost_breakdown: BTreeMap::new(),
|
||||
tier_index,
|
||||
tier_info,
|
||||
missing_required: Vec::new(),
|
||||
error: Some("negative_cost".to_string()),
|
||||
});
|
||||
}
|
||||
|
||||
let mut breakdown = BTreeMap::new();
|
||||
for (key, value) in &resolved {
|
||||
if key.ends_with("_cost") {
|
||||
if let Some(number) = as_f64(value) {
|
||||
breakdown.insert(key.clone(), quantize_cost(number));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(FormulaEvaluationResult {
|
||||
status: FormulaEvaluationStatus::Complete,
|
||||
cost: quantize_cost(cost),
|
||||
resolved_dimensions: dims,
|
||||
resolved_variables: resolved,
|
||||
cost_breakdown: breakdown,
|
||||
tier_index,
|
||||
tier_info,
|
||||
missing_required: Vec::new(),
|
||||
error: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn extract_variable_names(expression: &str) -> Result<BTreeSet<String>, UnsafeExpressionError> {
|
||||
let tokens = tokenize(expression)?;
|
||||
let mut names = BTreeSet::new();
|
||||
for index in 0..tokens.len() {
|
||||
if let Token::Identifier(name) = &tokens[index] {
|
||||
if matches!(tokens.get(index + 1), Some(Token::LeftParen)) {
|
||||
continue;
|
||||
}
|
||||
names.insert(name.clone());
|
||||
}
|
||||
}
|
||||
Ok(names)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum ComputedStatus {
|
||||
Ok,
|
||||
Pending,
|
||||
MissingRequired,
|
||||
Defaulted,
|
||||
}
|
||||
|
||||
fn try_resolve_computed(
|
||||
mapping: &serde_json::Value,
|
||||
dims: &BTreeMap<String, serde_json::Value>,
|
||||
resolved: &BTreeMap<String, serde_json::Value>,
|
||||
) -> Result<(ComputedStatus, serde_json::Value), ExpressionEvaluationError> {
|
||||
let required = mapping
|
||||
.get("required")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
let default = mapping
|
||||
.get("default")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| serde_json::json!(0));
|
||||
let expr = mapping
|
||||
.get("expression")
|
||||
.or_else(|| mapping.get("transform_expression"))
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or_default();
|
||||
if expr.is_empty() {
|
||||
return Ok((
|
||||
if required {
|
||||
ComputedStatus::MissingRequired
|
||||
} else {
|
||||
ComputedStatus::Defaulted
|
||||
},
|
||||
default,
|
||||
));
|
||||
}
|
||||
|
||||
let needed = extract_variable_names(expr)
|
||||
.map_err(|err| ExpressionEvaluationError::Failed(err.to_string()))?;
|
||||
let env: BTreeMap<String, serde_json::Value> = dims
|
||||
.iter()
|
||||
.map(|(k, v)| (k.clone(), v.clone()))
|
||||
.chain(resolved.iter().map(|(k, v)| (k.clone(), v.clone())))
|
||||
.collect();
|
||||
if needed.iter().any(|name| !env.contains_key(name)) {
|
||||
return Ok((
|
||||
if required {
|
||||
ComputedStatus::Pending
|
||||
} else {
|
||||
ComputedStatus::Defaulted
|
||||
},
|
||||
default,
|
||||
));
|
||||
}
|
||||
|
||||
match evaluate_expression(expr, &env) {
|
||||
Ok(value) => Ok((ComputedStatus::Ok, serde_json::json!(value))),
|
||||
Err(err) if required => Err(ExpressionEvaluationError::Failed(err.to_string())),
|
||||
Err(_) => Ok((ComputedStatus::Defaulted, default)),
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_mapping(
|
||||
var_name: &str,
|
||||
mapping: &serde_json::Value,
|
||||
dims: &BTreeMap<String, serde_json::Value>,
|
||||
) -> MappingResolution {
|
||||
let source = mapping
|
||||
.get("source")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("constant")
|
||||
.to_ascii_lowercase();
|
||||
|
||||
match source.as_str() {
|
||||
"constant" => Ok((
|
||||
mapping
|
||||
.get("default")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| serde_json::json!(0)),
|
||||
false,
|
||||
None,
|
||||
)),
|
||||
"dimension" => resolve_dimension(var_name, mapping, dims),
|
||||
"tiered" => resolve_tiered(var_name, mapping, dims),
|
||||
"matrix" => resolve_matrix(var_name, mapping, dims),
|
||||
_ => Ok((
|
||||
mapping
|
||||
.get("default")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| serde_json::json!(0)),
|
||||
false,
|
||||
None,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_dimension(
|
||||
var_name: &str,
|
||||
mapping: &serde_json::Value,
|
||||
dims: &BTreeMap<String, serde_json::Value>,
|
||||
) -> MappingResolution {
|
||||
let key = mapping
|
||||
.get("key")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or(var_name);
|
||||
let required = mapping
|
||||
.get("required")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
let allow_zero = mapping
|
||||
.get("allow_zero")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(true);
|
||||
let default = mapping
|
||||
.get("default")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| serde_json::json!(0));
|
||||
|
||||
let Some(value) = dims.get(key).cloned() else {
|
||||
return Ok((default, required, None));
|
||||
};
|
||||
if !allow_zero && is_zero_like(&value) {
|
||||
return Ok((default, required, None));
|
||||
}
|
||||
Ok((value, false, None))
|
||||
}
|
||||
|
||||
fn resolve_tiered(
|
||||
_var_name: &str,
|
||||
mapping: &serde_json::Value,
|
||||
dims: &BTreeMap<String, serde_json::Value>,
|
||||
) -> MappingResolution {
|
||||
let tier_key = mapping
|
||||
.get("tier_key")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("total_input_context");
|
||||
let default = mapping
|
||||
.get("default")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| serde_json::json!(0));
|
||||
let Some(tiers) = mapping.get("tiers").and_then(|v| v.as_array()) else {
|
||||
return Ok((default, false, None));
|
||||
};
|
||||
|
||||
let tier_value = dims.get(tier_key).and_then(as_f64).unwrap_or(0.0);
|
||||
let ttl_key = mapping.get("ttl_key").and_then(|v| v.as_str());
|
||||
let ttl_value_key = mapping.get("ttl_value_key").and_then(|v| v.as_str());
|
||||
let cache_ttl_minutes = ttl_key.and_then(|key| dims.get(key)).and_then(as_f64);
|
||||
|
||||
let mut matched_index = None;
|
||||
let mut matched_tier = None;
|
||||
for (index, tier) in tiers.iter().enumerate() {
|
||||
let up_to = tier.get("up_to").and_then(as_f64);
|
||||
if up_to.is_none() || tier_value <= up_to.unwrap_or_default() {
|
||||
matched_index = Some(index as i64);
|
||||
matched_tier = Some(tier.clone());
|
||||
break;
|
||||
}
|
||||
}
|
||||
let Some(tier) = matched_tier.or_else(|| tiers.last().cloned()) else {
|
||||
return Ok((default, false, None));
|
||||
};
|
||||
|
||||
if let (Some(cache_ttl_minutes), Some(ttl_value_key)) = (cache_ttl_minutes, ttl_value_key) {
|
||||
if let Some(ttl_pricing) = tier.get("cache_ttl_pricing").and_then(|v| v.as_array()) {
|
||||
for ttl_entry in ttl_pricing {
|
||||
let ttl_limit = ttl_entry
|
||||
.get("ttl_minutes")
|
||||
.and_then(as_f64)
|
||||
.unwrap_or_default();
|
||||
if cache_ttl_minutes <= ttl_limit {
|
||||
if let Some(value) = ttl_entry.get(ttl_value_key) {
|
||||
return Ok((
|
||||
value.clone(),
|
||||
false,
|
||||
Some((matched_index.unwrap_or(0), tier)),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(value) = ttl_pricing.last().and_then(|v| v.get(ttl_value_key)) {
|
||||
return Ok((
|
||||
value.clone(),
|
||||
false,
|
||||
Some((matched_index.unwrap_or(0), tier)),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok((
|
||||
tier.get("value").cloned().unwrap_or(default),
|
||||
false,
|
||||
Some((matched_index.unwrap_or(0), tier)),
|
||||
))
|
||||
}
|
||||
|
||||
fn resolve_matrix(
|
||||
var_name: &str,
|
||||
mapping: &serde_json::Value,
|
||||
dims: &BTreeMap<String, serde_json::Value>,
|
||||
) -> MappingResolution {
|
||||
let key = mapping
|
||||
.get("key")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or(var_name);
|
||||
let default = mapping
|
||||
.get("default")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| serde_json::json!(0));
|
||||
let Some(value) = dims.get(key) else {
|
||||
let required = mapping
|
||||
.get("required")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
return Ok((default, required, None));
|
||||
};
|
||||
let Some(entries) = mapping.get("entries").and_then(|v| v.as_object()) else {
|
||||
return Ok((default, false, None));
|
||||
};
|
||||
let lookup_key = match value {
|
||||
serde_json::Value::String(text) => text.clone(),
|
||||
_ => value.to_string(),
|
||||
};
|
||||
Ok((
|
||||
entries.get(&lookup_key).cloned().unwrap_or(default),
|
||||
false,
|
||||
None,
|
||||
))
|
||||
}
|
||||
|
||||
fn is_zero_like(value: &serde_json::Value) -> bool {
|
||||
as_f64(value).map(|number| number == 0.0).unwrap_or(false)
|
||||
}
|
||||
|
||||
fn as_f64(value: &serde_json::Value) -> Option<f64> {
|
||||
value.as_f64().or_else(|| value.as_i64().map(|v| v as f64))
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
enum Token {
|
||||
Number(f64),
|
||||
Identifier(String),
|
||||
Plus,
|
||||
Minus,
|
||||
Star,
|
||||
DoubleStar,
|
||||
Slash,
|
||||
DoubleSlash,
|
||||
Percent,
|
||||
LeftParen,
|
||||
RightParen,
|
||||
Comma,
|
||||
}
|
||||
|
||||
fn tokenize(input: &str) -> Result<Vec<Token>, UnsafeExpressionError> {
|
||||
let chars: Vec<char> = input.chars().collect();
|
||||
let mut tokens = Vec::new();
|
||||
let mut index = 0;
|
||||
while index < chars.len() {
|
||||
match chars[index] {
|
||||
ch if ch.is_ascii_whitespace() => index += 1,
|
||||
ch if ch.is_ascii_digit() || ch == '.' => {
|
||||
let start = index;
|
||||
index += 1;
|
||||
while index < chars.len() && (chars[index].is_ascii_digit() || chars[index] == '.')
|
||||
{
|
||||
index += 1;
|
||||
}
|
||||
let text: String = chars[start..index].iter().collect();
|
||||
let number = text.parse::<f64>().map_err(|_| {
|
||||
UnsafeExpressionError::Unsupported(format!("invalid numeric literal: {text}"))
|
||||
})?;
|
||||
tokens.push(Token::Number(number));
|
||||
}
|
||||
ch if ch.is_ascii_alphabetic() || ch == '_' => {
|
||||
let start = index;
|
||||
index += 1;
|
||||
while index < chars.len()
|
||||
&& (chars[index].is_ascii_alphanumeric() || chars[index] == '_')
|
||||
{
|
||||
index += 1;
|
||||
}
|
||||
tokens.push(Token::Identifier(chars[start..index].iter().collect()));
|
||||
}
|
||||
'+' => {
|
||||
tokens.push(Token::Plus);
|
||||
index += 1;
|
||||
}
|
||||
'-' => {
|
||||
tokens.push(Token::Minus);
|
||||
index += 1;
|
||||
}
|
||||
'*' => {
|
||||
if matches!(chars.get(index + 1), Some('*')) {
|
||||
tokens.push(Token::DoubleStar);
|
||||
index += 2;
|
||||
} else {
|
||||
tokens.push(Token::Star);
|
||||
index += 1;
|
||||
}
|
||||
}
|
||||
'/' => {
|
||||
if matches!(chars.get(index + 1), Some('/')) {
|
||||
tokens.push(Token::DoubleSlash);
|
||||
index += 2;
|
||||
} else {
|
||||
tokens.push(Token::Slash);
|
||||
index += 1;
|
||||
}
|
||||
}
|
||||
'%' => {
|
||||
tokens.push(Token::Percent);
|
||||
index += 1;
|
||||
}
|
||||
'(' => {
|
||||
tokens.push(Token::LeftParen);
|
||||
index += 1;
|
||||
}
|
||||
')' => {
|
||||
tokens.push(Token::RightParen);
|
||||
index += 1;
|
||||
}
|
||||
',' => {
|
||||
tokens.push(Token::Comma);
|
||||
index += 1;
|
||||
}
|
||||
other => {
|
||||
return Err(UnsafeExpressionError::Unsupported(format!(
|
||||
"unsupported token: {other}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(tokens)
|
||||
}
|
||||
|
||||
fn evaluate_expression(
|
||||
expression: &str,
|
||||
variables: &BTreeMap<String, serde_json::Value>,
|
||||
) -> Result<f64, UnsafeExpressionError> {
|
||||
let tokens = tokenize(expression)?;
|
||||
let mut parser = Parser {
|
||||
tokens: &tokens,
|
||||
index: 0,
|
||||
variables,
|
||||
};
|
||||
let value = parser.parse_expression()?;
|
||||
if parser.index != tokens.len() {
|
||||
return Err(UnsafeExpressionError::Unsupported(
|
||||
"unexpected trailing tokens".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
struct Parser<'a> {
|
||||
tokens: &'a [Token],
|
||||
index: usize,
|
||||
variables: &'a BTreeMap<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
impl<'a> Parser<'a> {
|
||||
fn parse_expression(&mut self) -> Result<f64, UnsafeExpressionError> {
|
||||
let mut left = self.parse_term()?;
|
||||
loop {
|
||||
match self.peek() {
|
||||
Some(Token::Plus) => {
|
||||
self.index += 1;
|
||||
left += self.parse_term()?;
|
||||
}
|
||||
Some(Token::Minus) => {
|
||||
self.index += 1;
|
||||
left -= self.parse_term()?;
|
||||
}
|
||||
_ => break,
|
||||
}
|
||||
}
|
||||
Ok(left)
|
||||
}
|
||||
|
||||
fn parse_term(&mut self) -> Result<f64, UnsafeExpressionError> {
|
||||
let mut left = self.parse_power()?;
|
||||
loop {
|
||||
match self.peek() {
|
||||
Some(Token::Star) => {
|
||||
self.index += 1;
|
||||
left *= self.parse_power()?;
|
||||
}
|
||||
Some(Token::Slash) => {
|
||||
self.index += 1;
|
||||
left /= self.parse_power()?;
|
||||
}
|
||||
Some(Token::DoubleSlash) => {
|
||||
self.index += 1;
|
||||
left = (left / self.parse_power()?).floor();
|
||||
}
|
||||
Some(Token::Percent) => {
|
||||
self.index += 1;
|
||||
left %= self.parse_power()?;
|
||||
}
|
||||
_ => break,
|
||||
}
|
||||
}
|
||||
Ok(left)
|
||||
}
|
||||
|
||||
fn parse_power(&mut self) -> Result<f64, UnsafeExpressionError> {
|
||||
let left = self.parse_unary()?;
|
||||
if matches!(self.peek(), Some(Token::DoubleStar)) {
|
||||
self.index += 1;
|
||||
let right = self.parse_power()?;
|
||||
return Ok(left.powf(right));
|
||||
}
|
||||
Ok(left)
|
||||
}
|
||||
|
||||
fn parse_unary(&mut self) -> Result<f64, UnsafeExpressionError> {
|
||||
match self.peek() {
|
||||
Some(Token::Plus) => {
|
||||
self.index += 1;
|
||||
self.parse_unary()
|
||||
}
|
||||
Some(Token::Minus) => {
|
||||
self.index += 1;
|
||||
Ok(-self.parse_unary()?)
|
||||
}
|
||||
_ => self.parse_primary(),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_primary(&mut self) -> Result<f64, UnsafeExpressionError> {
|
||||
match self.next() {
|
||||
Some(Token::Number(value)) => Ok(*value),
|
||||
Some(Token::Identifier(name)) => {
|
||||
if matches!(self.peek(), Some(Token::LeftParen)) {
|
||||
self.index += 1;
|
||||
let args = self.parse_arguments()?;
|
||||
evaluate_function(name, &args)
|
||||
} else {
|
||||
let value = self.variables.get(name).and_then(as_f64).ok_or_else(|| {
|
||||
UnsafeExpressionError::Unsupported(format!("unknown variable: {name}"))
|
||||
})?;
|
||||
Ok(value)
|
||||
}
|
||||
}
|
||||
Some(Token::LeftParen) => {
|
||||
let value = self.parse_expression()?;
|
||||
match self.next() {
|
||||
Some(Token::RightParen) => Ok(value),
|
||||
_ => Err(UnsafeExpressionError::Unsupported(
|
||||
"missing closing ')'".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
other => Err(UnsafeExpressionError::Unsupported(format!(
|
||||
"unexpected token: {other:?}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_arguments(&mut self) -> Result<Vec<f64>, UnsafeExpressionError> {
|
||||
let mut args = Vec::new();
|
||||
if matches!(self.peek(), Some(Token::RightParen)) {
|
||||
self.index += 1;
|
||||
return Ok(args);
|
||||
}
|
||||
loop {
|
||||
args.push(self.parse_expression()?);
|
||||
match self.next() {
|
||||
Some(Token::Comma) => {}
|
||||
Some(Token::RightParen) => break,
|
||||
other => {
|
||||
return Err(UnsafeExpressionError::Unsupported(format!(
|
||||
"unexpected token in function call: {other:?}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(args)
|
||||
}
|
||||
|
||||
fn peek(&self) -> Option<&'a Token> {
|
||||
self.tokens.get(self.index)
|
||||
}
|
||||
|
||||
fn next(&mut self) -> Option<&'a Token> {
|
||||
let token = self.tokens.get(self.index);
|
||||
if token.is_some() {
|
||||
self.index += 1;
|
||||
}
|
||||
token
|
||||
}
|
||||
}
|
||||
|
||||
fn evaluate_function(name: &str, args: &[f64]) -> Result<f64, UnsafeExpressionError> {
|
||||
match name {
|
||||
"min" => args.iter().copied().reduce(f64::min).ok_or_else(|| {
|
||||
UnsafeExpressionError::Unsupported("min() requires arguments".to_string())
|
||||
}),
|
||||
"max" => args.iter().copied().reduce(f64::max).ok_or_else(|| {
|
||||
UnsafeExpressionError::Unsupported("max() requires arguments".to_string())
|
||||
}),
|
||||
"abs" => args.first().copied().map(f64::abs).ok_or_else(|| {
|
||||
UnsafeExpressionError::Unsupported("abs() requires one argument".to_string())
|
||||
}),
|
||||
"round" => {
|
||||
let Some(value) = args.first().copied() else {
|
||||
return Err(UnsafeExpressionError::Unsupported(
|
||||
"round() requires at least one argument".to_string(),
|
||||
));
|
||||
};
|
||||
let digits = args.get(1).copied().unwrap_or(0.0) as i32;
|
||||
let factor = 10_f64.powi(digits);
|
||||
Ok((value * factor).round() / factor)
|
||||
}
|
||||
"int" => args
|
||||
.first()
|
||||
.copied()
|
||||
.map(|value| value.trunc())
|
||||
.ok_or_else(|| {
|
||||
UnsafeExpressionError::Unsupported("int() requires one argument".to_string())
|
||||
}),
|
||||
"float" => args.first().copied().ok_or_else(|| {
|
||||
UnsafeExpressionError::Unsupported("float() requires one argument".to_string())
|
||||
}),
|
||||
_ => Err(UnsafeExpressionError::Unsupported(format!(
|
||||
"function not allowed: {name}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{extract_variable_names, FormulaEngine, FormulaEvaluationStatus};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
#[test]
|
||||
fn extracts_variable_names_without_functions() {
|
||||
let names = extract_variable_names("min(input_cost, output_cost) + tax")
|
||||
.expect("variables should parse");
|
||||
assert!(names.contains("input_cost"));
|
||||
assert!(names.contains("output_cost"));
|
||||
assert!(names.contains("tax"));
|
||||
assert!(!names.contains("min"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evaluates_formula_with_computed_variables() {
|
||||
let engine = FormulaEngine::new();
|
||||
let dimensions = BTreeMap::from([
|
||||
("input_tokens".to_string(), serde_json::json!(1000)),
|
||||
("output_tokens".to_string(), serde_json::json!(500)),
|
||||
]);
|
||||
let variables = BTreeMap::from([
|
||||
("input_price_per_1m".to_string(), serde_json::json!(3.0)),
|
||||
("output_price_per_1m".to_string(), serde_json::json!(15.0)),
|
||||
]);
|
||||
let mappings = BTreeMap::from([
|
||||
(
|
||||
"input_cost".to_string(),
|
||||
serde_json::json!({
|
||||
"source": "computed",
|
||||
"expression": "input_tokens * input_price_per_1m / 1000000"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"output_cost".to_string(),
|
||||
serde_json::json!({
|
||||
"source": "computed",
|
||||
"expression": "output_tokens * output_price_per_1m / 1000000"
|
||||
}),
|
||||
),
|
||||
]);
|
||||
|
||||
let result = engine
|
||||
.evaluate(
|
||||
"input_cost + output_cost",
|
||||
Some(&variables),
|
||||
Some(&dimensions),
|
||||
Some(&mappings),
|
||||
false,
|
||||
)
|
||||
.expect("formula should evaluate");
|
||||
|
||||
assert_eq!(result.status, FormulaEvaluationStatus::Complete);
|
||||
assert_eq!(result.cost, 0.0105);
|
||||
assert_eq!(result.cost_breakdown.get("input_cost"), Some(&0.003_f64));
|
||||
assert_eq!(result.cost_breakdown.get("output_cost"), Some(&0.0075_f64));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_tiered_mapping() {
|
||||
let engine = FormulaEngine::new();
|
||||
let dimensions =
|
||||
BTreeMap::from([("total_input_context".to_string(), serde_json::json!(2000))]);
|
||||
let mappings = BTreeMap::from([(
|
||||
"input_price_per_1m".to_string(),
|
||||
serde_json::json!({
|
||||
"source": "tiered",
|
||||
"tier_key": "total_input_context",
|
||||
"tiers": [
|
||||
{ "up_to": 1000, "value": 1.0 },
|
||||
{ "up_to": 4000, "value": 2.0 }
|
||||
],
|
||||
"default": 0.0
|
||||
}),
|
||||
)]);
|
||||
|
||||
let result = engine
|
||||
.evaluate(
|
||||
"input_price_per_1m",
|
||||
None,
|
||||
Some(&dimensions),
|
||||
Some(&mappings),
|
||||
false,
|
||||
)
|
||||
.expect("tiered mapping should evaluate");
|
||||
|
||||
assert_eq!(result.cost, 2.0);
|
||||
assert_eq!(result.tier_index, Some(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn incomplete_when_required_dimension_missing() {
|
||||
let engine = FormulaEngine::new();
|
||||
let mappings = BTreeMap::from([(
|
||||
"input_tokens".to_string(),
|
||||
serde_json::json!({
|
||||
"source": "dimension",
|
||||
"key": "input_tokens",
|
||||
"required": true
|
||||
}),
|
||||
)]);
|
||||
|
||||
let result = engine
|
||||
.evaluate("input_tokens", None, None, Some(&mappings), false)
|
||||
.expect("incomplete result should return ok");
|
||||
|
||||
assert_eq!(result.status, FormulaEvaluationStatus::Incomplete);
|
||||
assert_eq!(result.missing_required, vec!["input_tokens".to_string()]);
|
||||
}
|
||||
}
|
||||
27
crates/aether-billing/src/lib.rs
Normal file
27
crates/aether-billing/src/lib.rs
Normal file
@@ -0,0 +1,27 @@
|
||||
mod default_rule;
|
||||
mod formula_engine;
|
||||
mod models;
|
||||
mod precision;
|
||||
mod pricing;
|
||||
mod schema;
|
||||
mod service;
|
||||
mod token_normalization;
|
||||
mod usage_mapper;
|
||||
|
||||
pub use default_rule::{normalize_task_type, DefaultBillingRuleGenerator, VirtualBillingRule};
|
||||
pub use formula_engine::{
|
||||
extract_variable_names, BillingIncompleteError, ExpressionEvaluationError, FormulaEngine,
|
||||
FormulaEvaluationResult, FormulaEvaluationStatus, UnsafeExpressionError,
|
||||
};
|
||||
pub use models::{BillingDimension, BillingUnit, CostBreakdown, StandardizedUsage};
|
||||
pub use precision::{
|
||||
quantize_cost, quantize_display, quantize_value, BILLING_DISPLAY_PRECISION,
|
||||
BILLING_STORAGE_PRECISION,
|
||||
};
|
||||
pub use pricing::{BillingComputation, BillingModelPricingSnapshot, BillingUsageInput};
|
||||
pub use schema::{
|
||||
BillingSnapshot, BillingSnapshotStatus, CostResult, BILLING_SNAPSHOT_SCHEMA_VERSION,
|
||||
};
|
||||
pub use service::BillingService;
|
||||
pub use token_normalization::normalize_input_tokens_for_billing;
|
||||
pub use usage_mapper::{map_usage, map_usage_from_response, UsageMapper};
|
||||
140
crates/aether-billing/src/models.rs
Normal file
140
crates/aether-billing/src/models.rs
Normal file
@@ -0,0 +1,140 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum BillingUnit {
|
||||
Per1MTokens,
|
||||
Per1MTokensHour,
|
||||
PerRequest,
|
||||
Fixed,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct BillingDimension {
|
||||
pub name: String,
|
||||
pub usage_field: String,
|
||||
pub price_field: String,
|
||||
pub unit: BillingUnit,
|
||||
pub default_price: f64,
|
||||
}
|
||||
|
||||
impl BillingDimension {
|
||||
pub fn calculate(&self, usage_value: f64, price: f64) -> f64 {
|
||||
if usage_value <= 0.0 || price <= 0.0 {
|
||||
return 0.0;
|
||||
}
|
||||
match self.unit {
|
||||
BillingUnit::Per1MTokens | BillingUnit::Per1MTokensHour => {
|
||||
(usage_value / 1_000_000.0) * price
|
||||
}
|
||||
BillingUnit::PerRequest => usage_value * price,
|
||||
BillingUnit::Fixed => price,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, Default)]
|
||||
pub struct StandardizedUsage {
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cache_creation_tokens: i64,
|
||||
pub cache_read_tokens: i64,
|
||||
pub reasoning_tokens: i64,
|
||||
pub cache_storage_token_hours: f64,
|
||||
pub request_count: i64,
|
||||
pub dimensions: BTreeMap<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
impl StandardizedUsage {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
request_count: 1,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get(&self, field_name: &str) -> Option<serde_json::Value> {
|
||||
match field_name {
|
||||
"input_tokens" => Some(serde_json::json!(self.input_tokens)),
|
||||
"output_tokens" => Some(serde_json::json!(self.output_tokens)),
|
||||
"cache_creation_tokens" => Some(serde_json::json!(self.cache_creation_tokens)),
|
||||
"cache_read_tokens" => Some(serde_json::json!(self.cache_read_tokens)),
|
||||
"reasoning_tokens" => Some(serde_json::json!(self.reasoning_tokens)),
|
||||
"cache_storage_token_hours" => Some(serde_json::json!(self.cache_storage_token_hours)),
|
||||
"request_count" => Some(serde_json::json!(self.request_count)),
|
||||
"extra" | "dimensions" => Some(serde_json::json!(self.dimensions)),
|
||||
_ => self.dimensions.get(field_name).cloned(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set(&mut self, field_name: &str, value: impl Into<serde_json::Value>) {
|
||||
let value = value.into();
|
||||
match field_name {
|
||||
"input_tokens" => self.input_tokens = as_i64(&value, 0),
|
||||
"output_tokens" => self.output_tokens = as_i64(&value, 0),
|
||||
"cache_creation_tokens" => self.cache_creation_tokens = as_i64(&value, 0),
|
||||
"cache_read_tokens" => self.cache_read_tokens = as_i64(&value, 0),
|
||||
"reasoning_tokens" => self.reasoning_tokens = as_i64(&value, 0),
|
||||
"cache_storage_token_hours" => self.cache_storage_token_hours = as_f64(&value, 0.0),
|
||||
"request_count" => self.request_count = as_i64(&value, 0),
|
||||
"extra" | "dimensions" => {
|
||||
self.dimensions = match value {
|
||||
serde_json::Value::Object(map) => map.into_iter().collect(),
|
||||
_ => BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
self.dimensions.insert(field_name.to_string(), value);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, Default)]
|
||||
pub struct CostBreakdown {
|
||||
pub costs: BTreeMap<String, f64>,
|
||||
pub total_cost: f64,
|
||||
pub tier_index: Option<i64>,
|
||||
pub effective_prices: BTreeMap<String, f64>,
|
||||
}
|
||||
|
||||
fn as_i64(value: &serde_json::Value, default: i64) -> i64 {
|
||||
value
|
||||
.as_i64()
|
||||
.or_else(|| value.as_u64().and_then(|v| i64::try_from(v).ok()))
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
fn as_f64(value: &serde_json::Value, default: f64) -> f64 {
|
||||
value.as_f64().unwrap_or(default)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{BillingDimension, BillingUnit, StandardizedUsage};
|
||||
|
||||
#[test]
|
||||
fn dimension_calculates_per_million_tokens() {
|
||||
let dimension = BillingDimension {
|
||||
name: "input".to_string(),
|
||||
usage_field: "input_tokens".to_string(),
|
||||
price_field: "input_price_per_1m".to_string(),
|
||||
unit: BillingUnit::Per1MTokens,
|
||||
default_price: 0.0,
|
||||
};
|
||||
assert_eq!(dimension.calculate(500_000.0, 2.0), 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn standardized_usage_reads_and_writes_known_and_extra_fields() {
|
||||
let mut usage = StandardizedUsage::new();
|
||||
usage.set("input_tokens", 10);
|
||||
usage.set("custom_dimension", "value");
|
||||
|
||||
assert_eq!(usage.get("input_tokens"), Some(serde_json::json!(10)));
|
||||
assert_eq!(
|
||||
usage.get("custom_dimension"),
|
||||
Some(serde_json::json!("value"))
|
||||
);
|
||||
}
|
||||
}
|
||||
33
crates/aether-billing/src/precision.rs
Normal file
33
crates/aether-billing/src/precision.rs
Normal file
@@ -0,0 +1,33 @@
|
||||
pub const BILLING_STORAGE_PRECISION: u32 = 8;
|
||||
pub const BILLING_DISPLAY_PRECISION: u32 = 6;
|
||||
|
||||
pub fn quantize_value(value: f64, precision: u32) -> f64 {
|
||||
if !value.is_finite() {
|
||||
return value;
|
||||
}
|
||||
let factor = 10_f64.powi(precision as i32);
|
||||
(value * factor).round() / factor
|
||||
}
|
||||
|
||||
pub fn quantize_cost(value: f64) -> f64 {
|
||||
quantize_value(value, BILLING_STORAGE_PRECISION)
|
||||
}
|
||||
|
||||
pub fn quantize_display(value: f64) -> f64 {
|
||||
quantize_value(value, BILLING_DISPLAY_PRECISION)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{quantize_cost, quantize_display};
|
||||
|
||||
#[test]
|
||||
fn quantizes_cost_to_storage_precision() {
|
||||
assert_eq!(quantize_cost(1.234567891), 1.23456789);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quantizes_display_to_display_precision() {
|
||||
assert_eq!(quantize_display(1.23456789), 1.234568);
|
||||
}
|
||||
}
|
||||
94
crates/aether-billing/src/pricing.rs
Normal file
94
crates/aether-billing/src/pricing.rs
Normal file
@@ -0,0 +1,94 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct BillingModelPricingSnapshot {
|
||||
pub provider_id: String,
|
||||
pub provider_billing_type: Option<String>,
|
||||
pub provider_api_key_id: Option<String>,
|
||||
pub provider_api_key_rate_multipliers: Option<Value>,
|
||||
pub provider_api_key_cache_ttl_minutes: Option<i64>,
|
||||
pub global_model_id: String,
|
||||
pub global_model_name: String,
|
||||
pub global_model_config: Option<Value>,
|
||||
pub default_price_per_request: Option<f64>,
|
||||
pub default_tiered_pricing: Option<Value>,
|
||||
pub model_id: Option<String>,
|
||||
pub model_provider_model_name: Option<String>,
|
||||
pub model_config: Option<Value>,
|
||||
pub model_price_per_request: Option<f64>,
|
||||
pub model_tiered_pricing: Option<Value>,
|
||||
}
|
||||
|
||||
impl BillingModelPricingSnapshot {
|
||||
pub fn effective_tiered_pricing(&self) -> Option<&Value> {
|
||||
self.model_tiered_pricing
|
||||
.as_ref()
|
||||
.or(self.default_tiered_pricing.as_ref())
|
||||
}
|
||||
|
||||
pub fn effective_price_per_request(&self) -> Option<f64> {
|
||||
self.model_price_per_request
|
||||
.or(self.default_price_per_request)
|
||||
}
|
||||
|
||||
pub fn is_free_tier(&self) -> bool {
|
||||
self.provider_billing_type
|
||||
.as_deref()
|
||||
.map(|value| value.eq_ignore_ascii_case("free_tier"))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
pub fn rate_multiplier_for_api_format(&self, api_format: Option<&str>) -> f64 {
|
||||
let Some(api_format) = api_format.map(str::trim).filter(|value| !value.is_empty()) else {
|
||||
return 1.0;
|
||||
};
|
||||
let normalized = api_format.to_ascii_lowercase();
|
||||
let Some(mapping) = self
|
||||
.provider_api_key_rate_multipliers
|
||||
.as_ref()
|
||||
.and_then(Value::as_object)
|
||||
else {
|
||||
return 1.0;
|
||||
};
|
||||
mapping
|
||||
.get(&normalized)
|
||||
.and_then(|value| value.as_f64())
|
||||
.unwrap_or(1.0)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct BillingUsageInput {
|
||||
pub task_type: String,
|
||||
pub api_format: Option<String>,
|
||||
pub request_count: i64,
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cache_creation_tokens: i64,
|
||||
pub cache_read_tokens: i64,
|
||||
pub cache_ttl_minutes: Option<i64>,
|
||||
}
|
||||
|
||||
impl BillingUsageInput {
|
||||
pub fn new(task_type: impl Into<String>) -> Self {
|
||||
Self {
|
||||
task_type: task_type.into(),
|
||||
api_format: None,
|
||||
request_count: 1,
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
cache_creation_tokens: 0,
|
||||
cache_read_tokens: 0,
|
||||
cache_ttl_minutes: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct BillingComputation {
|
||||
pub cost_result: crate::CostResult,
|
||||
pub actual_total_cost: f64,
|
||||
pub rate_multiplier: f64,
|
||||
pub is_free_tier: bool,
|
||||
}
|
||||
110
crates/aether-billing/src/schema.rs
Normal file
110
crates/aether-billing/src/schema.rs
Normal file
@@ -0,0 +1,110 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
pub const BILLING_SNAPSHOT_SCHEMA_VERSION: &str = "2.0";
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum BillingSnapshotStatus {
|
||||
Complete,
|
||||
Incomplete,
|
||||
NoRule,
|
||||
Legacy,
|
||||
}
|
||||
|
||||
impl BillingSnapshotStatus {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Complete => "complete",
|
||||
Self::Incomplete => "incomplete",
|
||||
Self::NoRule => "no_rule",
|
||||
Self::Legacy => "legacy",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for BillingSnapshotStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str(self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct BillingSnapshot {
|
||||
pub schema_version: String,
|
||||
pub rule_id: Option<String>,
|
||||
pub rule_name: Option<String>,
|
||||
pub scope: Option<String>,
|
||||
pub expression: Option<String>,
|
||||
pub resolved_dimensions: BTreeMap<String, serde_json::Value>,
|
||||
pub resolved_variables: BTreeMap<String, serde_json::Value>,
|
||||
pub cost_breakdown: BTreeMap<String, f64>,
|
||||
pub total_cost: f64,
|
||||
pub tier_index: Option<i64>,
|
||||
pub tier_info: Option<serde_json::Value>,
|
||||
pub missing_required: Vec<String>,
|
||||
pub status: BillingSnapshotStatus,
|
||||
pub calculated_at: String,
|
||||
pub engine_version: String,
|
||||
}
|
||||
|
||||
impl Default for BillingSnapshot {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
schema_version: BILLING_SNAPSHOT_SCHEMA_VERSION.to_string(),
|
||||
rule_id: None,
|
||||
rule_name: None,
|
||||
scope: None,
|
||||
expression: None,
|
||||
resolved_dimensions: BTreeMap::new(),
|
||||
resolved_variables: BTreeMap::new(),
|
||||
cost_breakdown: BTreeMap::new(),
|
||||
total_cost: 0.0,
|
||||
tier_index: None,
|
||||
tier_info: None,
|
||||
missing_required: Vec::new(),
|
||||
status: BillingSnapshotStatus::NoRule,
|
||||
calculated_at: String::new(),
|
||||
engine_version: "2.0".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl BillingSnapshot {
|
||||
pub fn dimensions_used(&self) -> &BTreeMap<String, serde_json::Value> {
|
||||
&self.resolved_dimensions
|
||||
}
|
||||
|
||||
pub fn cost(&self) -> f64 {
|
||||
self.total_cost
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct CostResult {
|
||||
pub cost: f64,
|
||||
pub status: BillingSnapshotStatus,
|
||||
pub snapshot: BillingSnapshot,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{BillingSnapshot, BillingSnapshotStatus};
|
||||
|
||||
#[test]
|
||||
fn snapshot_aliases_dimensions_and_cost() {
|
||||
let mut snapshot = BillingSnapshot {
|
||||
total_cost: 1.25,
|
||||
status: BillingSnapshotStatus::Complete,
|
||||
..BillingSnapshot::default()
|
||||
};
|
||||
snapshot
|
||||
.resolved_dimensions
|
||||
.insert("input_tokens".to_string(), serde_json::json!(10));
|
||||
|
||||
assert_eq!(snapshot.cost(), 1.25);
|
||||
assert_eq!(
|
||||
snapshot.dimensions_used().get("input_tokens"),
|
||||
Some(&serde_json::json!(10))
|
||||
);
|
||||
}
|
||||
}
|
||||
237
crates/aether-billing/src/service.rs
Normal file
237
crates/aether-billing/src/service.rs
Normal file
@@ -0,0 +1,237 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use crate::default_rule::{normalize_task_type, DefaultBillingRuleGenerator};
|
||||
use crate::precision::quantize_cost;
|
||||
use crate::pricing::{BillingComputation, BillingModelPricingSnapshot, BillingUsageInput};
|
||||
use crate::schema::{
|
||||
BillingSnapshot, BillingSnapshotStatus, CostResult, BILLING_SNAPSHOT_SCHEMA_VERSION,
|
||||
};
|
||||
use crate::{
|
||||
normalize_input_tokens_for_billing, ExpressionEvaluationError, FormulaEngine,
|
||||
FormulaEvaluationStatus,
|
||||
};
|
||||
|
||||
pub struct BillingService {
|
||||
engine: FormulaEngine,
|
||||
}
|
||||
|
||||
impl BillingService {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
engine: FormulaEngine::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn calculate(
|
||||
&self,
|
||||
pricing: &BillingModelPricingSnapshot,
|
||||
input: &BillingUsageInput,
|
||||
) -> Result<BillingComputation, ExpressionEvaluationError> {
|
||||
let Some(rule) =
|
||||
DefaultBillingRuleGenerator::generate_for_pricing(pricing, &input.task_type)
|
||||
else {
|
||||
return Ok(BillingComputation {
|
||||
cost_result: CostResult {
|
||||
cost: 0.0,
|
||||
status: BillingSnapshotStatus::NoRule,
|
||||
snapshot: BillingSnapshot {
|
||||
schema_version: BILLING_SNAPSHOT_SCHEMA_VERSION.to_string(),
|
||||
rule_id: None,
|
||||
rule_name: None,
|
||||
scope: None,
|
||||
expression: None,
|
||||
resolved_dimensions: build_dimensions(input),
|
||||
resolved_variables: BTreeMap::new(),
|
||||
cost_breakdown: BTreeMap::new(),
|
||||
total_cost: 0.0,
|
||||
tier_index: None,
|
||||
tier_info: None,
|
||||
missing_required: Vec::new(),
|
||||
status: BillingSnapshotStatus::NoRule,
|
||||
calculated_at: now_marker(),
|
||||
engine_version: "2.0".to_string(),
|
||||
},
|
||||
},
|
||||
actual_total_cost: 0.0,
|
||||
rate_multiplier: pricing
|
||||
.rate_multiplier_for_api_format(input.api_format.as_deref()),
|
||||
is_free_tier: pricing.is_free_tier(),
|
||||
});
|
||||
};
|
||||
|
||||
let dims = build_dimensions(input);
|
||||
let result = self.engine.evaluate(
|
||||
&rule.expression,
|
||||
Some(&rule.variables),
|
||||
Some(&dims),
|
||||
Some(&rule.dimension_mappings),
|
||||
false,
|
||||
)?;
|
||||
|
||||
let status = match result.status {
|
||||
FormulaEvaluationStatus::Complete => BillingSnapshotStatus::Complete,
|
||||
FormulaEvaluationStatus::Incomplete => BillingSnapshotStatus::Incomplete,
|
||||
};
|
||||
let total_cost = if matches!(status, BillingSnapshotStatus::Complete) {
|
||||
result.cost
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
let rate_multiplier = pricing.rate_multiplier_for_api_format(input.api_format.as_deref());
|
||||
let is_free_tier = pricing.is_free_tier();
|
||||
let actual_total_cost = if is_free_tier {
|
||||
0.0
|
||||
} else {
|
||||
quantize_cost(total_cost * rate_multiplier)
|
||||
};
|
||||
|
||||
Ok(BillingComputation {
|
||||
cost_result: CostResult {
|
||||
cost: total_cost,
|
||||
status,
|
||||
snapshot: BillingSnapshot {
|
||||
schema_version: BILLING_SNAPSHOT_SCHEMA_VERSION.to_string(),
|
||||
rule_id: Some(rule.id),
|
||||
rule_name: Some(rule.name),
|
||||
scope: Some(rule.scope),
|
||||
expression: Some(rule.expression),
|
||||
resolved_dimensions: result.resolved_dimensions,
|
||||
resolved_variables: result.resolved_variables,
|
||||
cost_breakdown: result.cost_breakdown,
|
||||
total_cost,
|
||||
tier_index: result.tier_index,
|
||||
tier_info: result.tier_info,
|
||||
missing_required: result.missing_required,
|
||||
status,
|
||||
calculated_at: now_marker(),
|
||||
engine_version: "2.0".to_string(),
|
||||
},
|
||||
},
|
||||
actual_total_cost,
|
||||
rate_multiplier,
|
||||
is_free_tier,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for BillingService {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
fn build_dimensions(input: &BillingUsageInput) -> BTreeMap<String, Value> {
|
||||
let normalized_input_tokens = normalize_input_tokens_for_billing(
|
||||
input.api_format.as_deref(),
|
||||
input.input_tokens,
|
||||
input.cache_read_tokens,
|
||||
);
|
||||
let total_input_context = input
|
||||
.input_tokens
|
||||
.saturating_add(input.cache_creation_tokens)
|
||||
.saturating_add(input.cache_read_tokens);
|
||||
|
||||
let mut out = BTreeMap::from([
|
||||
("input_tokens".to_string(), json!(normalized_input_tokens)),
|
||||
("output_tokens".to_string(), json!(input.output_tokens)),
|
||||
(
|
||||
"cache_creation_tokens".to_string(),
|
||||
json!(input.cache_creation_tokens),
|
||||
),
|
||||
(
|
||||
"cache_read_tokens".to_string(),
|
||||
json!(input.cache_read_tokens),
|
||||
),
|
||||
(
|
||||
"request_count".to_string(),
|
||||
json!(input.request_count.max(0)),
|
||||
),
|
||||
(
|
||||
"total_input_context".to_string(),
|
||||
json!(total_input_context),
|
||||
),
|
||||
(
|
||||
"effective_task_type".to_string(),
|
||||
json!(normalize_task_type(&input.task_type)),
|
||||
),
|
||||
]);
|
||||
|
||||
if let Some(cache_ttl_minutes) = input.cache_ttl_minutes {
|
||||
out.insert(
|
||||
"cache_ttl_minutes".to_string(),
|
||||
json!(cache_ttl_minutes.max(0)),
|
||||
);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn now_marker() -> String {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
.to_string()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::BillingService;
|
||||
use crate::{BillingModelPricingSnapshot, BillingSnapshotStatus, BillingUsageInput};
|
||||
|
||||
fn pricing() -> BillingModelPricingSnapshot {
|
||||
BillingModelPricingSnapshot {
|
||||
provider_id: "provider-1".to_string(),
|
||||
provider_billing_type: Some("pay_as_you_go".to_string()),
|
||||
provider_api_key_id: Some("key-1".to_string()),
|
||||
provider_api_key_rate_multipliers: Some(json!({"openai:chat": 0.5})),
|
||||
provider_api_key_cache_ttl_minutes: Some(60),
|
||||
global_model_id: "global-model-1".to_string(),
|
||||
global_model_name: "gpt-5".to_string(),
|
||||
global_model_config: None,
|
||||
default_price_per_request: Some(0.02),
|
||||
default_tiered_pricing: Some(json!({
|
||||
"tiers": [{
|
||||
"up_to": null,
|
||||
"input_price_per_1m": 3.0,
|
||||
"output_price_per_1m": 15.0,
|
||||
"cache_creation_price_per_1m": 3.75,
|
||||
"cache_read_price_per_1m": 0.30
|
||||
}]
|
||||
})),
|
||||
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: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn calculates_complete_snapshot_for_usage() {
|
||||
let result = BillingService::new()
|
||||
.calculate(
|
||||
&pricing(),
|
||||
&BillingUsageInput {
|
||||
task_type: "chat".to_string(),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
request_count: 1,
|
||||
input_tokens: 1_000,
|
||||
output_tokens: 500,
|
||||
cache_creation_tokens: 0,
|
||||
cache_read_tokens: 100,
|
||||
cache_ttl_minutes: Some(60),
|
||||
},
|
||||
)
|
||||
.expect("billing should calculate");
|
||||
|
||||
assert_eq!(result.cost_result.status, BillingSnapshotStatus::Complete);
|
||||
assert!(result.cost_result.cost > 0.0);
|
||||
assert!(result.actual_total_cost > 0.0);
|
||||
assert_eq!(result.rate_multiplier, 0.5);
|
||||
}
|
||||
}
|
||||
69
crates/aether-billing/src/token_normalization.rs
Normal file
69
crates/aether-billing/src/token_normalization.rs
Normal file
@@ -0,0 +1,69 @@
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum ApiFamily {
|
||||
OpenAi,
|
||||
Claude,
|
||||
Gemini,
|
||||
Unknown,
|
||||
}
|
||||
|
||||
fn parse_api_family(api_format: Option<&str>) -> ApiFamily {
|
||||
let Some(api_format) = api_format else {
|
||||
return ApiFamily::Unknown;
|
||||
};
|
||||
let family = api_format
|
||||
.split(':')
|
||||
.next()
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
match family.as_str() {
|
||||
"openai" => ApiFamily::OpenAi,
|
||||
"claude" | "anthropic" => ApiFamily::Claude,
|
||||
"gemini" | "google" => ApiFamily::Gemini,
|
||||
_ => ApiFamily::Unknown,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn normalize_input_tokens_for_billing(
|
||||
api_format: Option<&str>,
|
||||
input_tokens: i64,
|
||||
cache_read_tokens: i64,
|
||||
) -> i64 {
|
||||
if input_tokens <= 0 {
|
||||
return input_tokens.max(0);
|
||||
}
|
||||
if cache_read_tokens <= 0 {
|
||||
return input_tokens;
|
||||
}
|
||||
|
||||
match parse_api_family(api_format) {
|
||||
ApiFamily::Claude => input_tokens,
|
||||
ApiFamily::OpenAi | ApiFamily::Gemini => (input_tokens - cache_read_tokens).max(0),
|
||||
ApiFamily::Unknown => input_tokens,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::normalize_input_tokens_for_billing;
|
||||
|
||||
#[test]
|
||||
fn subtracts_cache_tokens_for_openai_and_gemini() {
|
||||
assert_eq!(
|
||||
normalize_input_tokens_for_billing(Some("openai:chat"), 100, 20),
|
||||
80
|
||||
);
|
||||
assert_eq!(
|
||||
normalize_input_tokens_for_billing(Some("gemini:chat"), 100, 20),
|
||||
80
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_input_tokens_for_claude() {
|
||||
assert_eq!(
|
||||
normalize_input_tokens_for_billing(Some("claude:chat"), 100, 20),
|
||||
100
|
||||
);
|
||||
}
|
||||
}
|
||||
190
crates/aether-billing/src/usage_mapper.rs
Normal file
190
crates/aether-billing/src/usage_mapper.rs
Normal file
@@ -0,0 +1,190 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::models::StandardizedUsage;
|
||||
|
||||
pub struct UsageMapper;
|
||||
|
||||
impl UsageMapper {
|
||||
pub fn map(
|
||||
raw_usage: &serde_json::Value,
|
||||
api_format: &str,
|
||||
extra_mapping: Option<&BTreeMap<String, String>>,
|
||||
) -> StandardizedUsage {
|
||||
if !raw_usage.is_object() {
|
||||
return StandardizedUsage::new();
|
||||
}
|
||||
|
||||
let mut usage = StandardizedUsage::new();
|
||||
let mut mapping = base_mapping(api_format);
|
||||
if let Some(extra_mapping) = extra_mapping {
|
||||
mapping.extend(extra_mapping.clone());
|
||||
}
|
||||
|
||||
for (source_path, target_field) in mapping {
|
||||
if let Some(value) = get_nested_value(raw_usage, &source_path) {
|
||||
usage.set(&target_field, value.clone());
|
||||
}
|
||||
}
|
||||
|
||||
usage
|
||||
}
|
||||
|
||||
pub fn map_from_response(response: &serde_json::Value, api_format: &str) -> StandardizedUsage {
|
||||
let family = api_family(api_format);
|
||||
let usage_value = if family == "gemini" {
|
||||
response
|
||||
.get("usageMetadata")
|
||||
.or_else(|| {
|
||||
response
|
||||
.get("candidates")
|
||||
.and_then(|v| v.get(0))
|
||||
.and_then(|v| v.get("usageMetadata"))
|
||||
})
|
||||
.unwrap_or(&serde_json::Value::Null)
|
||||
} else {
|
||||
response.get("usage").unwrap_or(&serde_json::Value::Null)
|
||||
};
|
||||
Self::map(usage_value, api_format, None)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn map_usage(raw_usage: &serde_json::Value, api_format: &str) -> StandardizedUsage {
|
||||
UsageMapper::map(raw_usage, api_format, None)
|
||||
}
|
||||
|
||||
pub fn map_usage_from_response(
|
||||
response: &serde_json::Value,
|
||||
api_format: &str,
|
||||
) -> StandardizedUsage {
|
||||
UsageMapper::map_from_response(response, api_format)
|
||||
}
|
||||
|
||||
fn api_family(api_format: &str) -> String {
|
||||
api_format
|
||||
.split(':')
|
||||
.next()
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn base_mapping(api_format: &str) -> BTreeMap<String, String> {
|
||||
let mut mapping = BTreeMap::new();
|
||||
match api_family(api_format).as_str() {
|
||||
"openai" => {
|
||||
mapping.insert("prompt_tokens".to_string(), "input_tokens".to_string());
|
||||
mapping.insert("completion_tokens".to_string(), "output_tokens".to_string());
|
||||
mapping.insert(
|
||||
"prompt_tokens_details.cached_tokens".to_string(),
|
||||
"cache_read_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"completion_tokens_details.reasoning_tokens".to_string(),
|
||||
"reasoning_tokens".to_string(),
|
||||
);
|
||||
}
|
||||
"gemini" => {
|
||||
mapping.insert("promptTokenCount".to_string(), "input_tokens".to_string());
|
||||
mapping.insert(
|
||||
"candidatesTokenCount".to_string(),
|
||||
"output_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"cachedContentTokenCount".to_string(),
|
||||
"cache_read_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"usageMetadata.promptTokenCount".to_string(),
|
||||
"input_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"usageMetadata.candidatesTokenCount".to_string(),
|
||||
"output_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"usageMetadata.cachedContentTokenCount".to_string(),
|
||||
"cache_read_tokens".to_string(),
|
||||
);
|
||||
}
|
||||
_ => {
|
||||
mapping.insert("input_tokens".to_string(), "input_tokens".to_string());
|
||||
mapping.insert("output_tokens".to_string(), "output_tokens".to_string());
|
||||
mapping.insert(
|
||||
"cache_creation_input_tokens".to_string(),
|
||||
"cache_creation_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"cache_read_input_tokens".to_string(),
|
||||
"cache_read_tokens".to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
mapping
|
||||
}
|
||||
|
||||
fn get_nested_value<'a>(value: &'a serde_json::Value, path: &str) -> Option<&'a serde_json::Value> {
|
||||
let mut current = value;
|
||||
for segment in path.split('.') {
|
||||
current = current.get(segment)?;
|
||||
}
|
||||
Some(current)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{map_usage, map_usage_from_response};
|
||||
|
||||
#[test]
|
||||
fn maps_openai_usage() {
|
||||
let usage = map_usage(
|
||||
&serde_json::json!({
|
||||
"prompt_tokens": 12,
|
||||
"completion_tokens": 8,
|
||||
"prompt_tokens_details": { "cached_tokens": 2 },
|
||||
"completion_tokens_details": { "reasoning_tokens": 3 }
|
||||
}),
|
||||
"openai:chat",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 12);
|
||||
assert_eq!(usage.output_tokens, 8);
|
||||
assert_eq!(usage.cache_read_tokens, 2);
|
||||
assert_eq!(usage.reasoning_tokens, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_claude_usage() {
|
||||
let usage = map_usage(
|
||||
&serde_json::json!({
|
||||
"input_tokens": 10,
|
||||
"output_tokens": 5,
|
||||
"cache_creation_input_tokens": 4,
|
||||
"cache_read_input_tokens": 1
|
||||
}),
|
||||
"claude:chat",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 10);
|
||||
assert_eq!(usage.output_tokens, 5);
|
||||
assert_eq!(usage.cache_creation_tokens, 4);
|
||||
assert_eq!(usage.cache_read_tokens, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_gemini_usage_from_response() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 14,
|
||||
"candidatesTokenCount": 6,
|
||||
"cachedContentTokenCount": 2
|
||||
}
|
||||
}),
|
||||
"gemini:chat",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 14);
|
||||
assert_eq!(usage.output_tokens, 6);
|
||||
assert_eq!(usage.cache_read_tokens, 2);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user