mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +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:
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()]);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user