use tau_proto::EstimatedApiCost;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct AgentCostSnapshot {
self_cost: EstimatedApiCost,
creator_subtree_cost: EstimatedApiCost,
}
impl AgentCostSnapshot {
pub(crate) fn new(self_cost: EstimatedApiCost, creator_subtree_cost: EstimatedApiCost) -> Self {
Self {
self_cost,
creator_subtree_cost,
}
}
}
#[derive(Clone, Default)]
pub(crate) struct AgentCostProjection {
costs: std::sync::Arc<
std::sync::Mutex<std::collections::HashMap<tau_proto::AgentId, AgentCostSnapshot>>,
>,
}
impl AgentCostProjection {
pub(crate) fn record(&self, agent_id: tau_proto::AgentId, costs: AgentCostSnapshot) {
crate::locked(&self.costs).insert(agent_id, costs);
}
pub(crate) fn clear(&self) {
crate::locked(&self.costs).clear();
}
pub(crate) fn snapshot(
&self,
) -> std::collections::HashMap<tau_proto::AgentId, AgentCostSnapshot> {
crate::locked(&self.costs).clone()
}
}
const PICODOLLARS_PER_CENT: u64 = 10_000_000_000;
const PICODOLLARS_PER_TENTH_DOLLAR: u64 = 100_000_000_000;
const PICODOLLARS_PER_DOLLAR: u64 = 1_000_000_000_000;
const PICODOLLARS_PER_THOUSAND_DOLLARS: u64 = 1_000 * PICODOLLARS_PER_DOLLAR;
const PICODOLLARS_PER_TENTH_MILLION_DOLLARS: u64 = 100_000 * PICODOLLARS_PER_DOLLAR;
const PICODOLLARS_PER_MILLION_DOLLARS: u64 = 1_000_000 * PICODOLLARS_PER_DOLLAR;
#[must_use]
pub(crate) fn format_compact(cost: EstimatedApiCost) -> String {
let picodollars = cost.as_picodollars();
let cents = rounded_units(picodollars, PICODOLLARS_PER_CENT);
if cents < 100 {
return format!("$.{cents:02}");
}
let tenths = rounded_units(picodollars, PICODOLLARS_PER_TENTH_DOLLAR);
if tenths < 100 {
return format!("${}.{}", tenths / 10, tenths % 10);
}
let dollars = rounded_units(picodollars, PICODOLLARS_PER_DOLLAR);
if dollars < 1_000 {
return format!("${dollars}");
}
let thousands = rounded_units(picodollars, PICODOLLARS_PER_THOUSAND_DOLLARS);
if thousands < 100 {
return format!("${thousands}k");
}
let tenths_of_millions = rounded_units(picodollars, PICODOLLARS_PER_TENTH_MILLION_DOLLARS);
if tenths_of_millions < 10 {
return format!("$.{tenths_of_millions}m");
}
let millions = rounded_units(picodollars, PICODOLLARS_PER_MILLION_DOLLARS);
format!("${millions}m")
}
#[must_use]
pub(crate) fn format_compact_pair(
self_cost: Option<EstimatedApiCost>,
creator_subtree_cost: Option<EstimatedApiCost>,
) -> String {
match (self_cost, creator_subtree_cost) {
(Some(self_cost), Some(creator_subtree_cost)) => format!(
"{}/{}",
format_compact(self_cost),
format_compact(creator_subtree_cost)
),
(Some(self_cost), None) => format!("{}/-", format_compact(self_cost)),
(None, _) => "-/-".to_owned(),
}
}
#[must_use]
pub(crate) fn format_snapshot(costs: Option<AgentCostSnapshot>) -> String {
let (self_cost, creator_subtree_cost) = costs
.map(|costs| (Some(costs.self_cost), Some(costs.creator_subtree_cost)))
.unwrap_or((None, None));
format_compact_pair(self_cost, creator_subtree_cost)
}
fn rounded_units(value: u64, unit: u64) -> u64 {
u64::try_from((u128::from(value) + u128::from(unit / 2)) / u128::from(unit)).unwrap_or(u64::MAX)
}
#[cfg(test)]
mod tests;