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:
fawney19
2026-03-31 19:19:04 +08:00
parent b5a0070023
commit ddf18fed9a
690 changed files with 235087 additions and 16301 deletions

View 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

View 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()
}

View 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()]);
}
}

View 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};

View 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"))
);
}
}

View 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);
}
}

View 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,
}

View 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))
);
}
}

View 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);
}
}

View 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
);
}
}

View 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);
}
}