use crate::budgets::{self, BudgetSpend, Scope, StepBill};
use std::path::PathBuf;
mod ceiling;
mod prices;
pub use ceiling::Ceiling;
pub use prices::{MICRO_PER_CENT, MICRO_PER_USD, Price, Prices};
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Cost {
pub micro_usd: u64,
pub unpriced_tokens: u64,
}
impl Cost {
pub fn usd(&self) -> String {
if self.micro_usd > 0 && self.micro_usd < MICRO_PER_CENT {
return "<$0.01".to_owned();
}
format!(
"${}.{:02}",
self.micro_usd / MICRO_PER_USD,
(self.micro_usd % MICRO_PER_USD) / MICRO_PER_CENT
)
}
fn add(&mut self, micro_usd: u64, unpriced: u64) {
self.micro_usd = self.micro_usd.saturating_add(micro_usd);
self.unpriced_tokens = self.unpriced_tokens.saturating_add(unpriced);
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Attribution {
Conversations(usize),
Workspace,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Note {
pub label: String,
pub hover: String,
}
impl Attribution {
pub fn note(&self) -> Option<Note> {
match self {
Self::Conversations(1) => None,
Self::Conversations(n) => Some(Note {
label: format!("over {n} conversations"),
hover: format!(
"{n} conversations name this ball in their goal. The figure \
sums all of them, sub-agents included."
),
}),
Self::Workspace => Some(Note {
label: "workspace-wide".to_owned(),
hover: "No conversation names this ball in its goal — it was picked up \
mid-conversation, and a pickup records only the workspace, never \
which conversation did it. So this is the whole workspace's spend: \
an upper bound on the ball, not the ball alone."
.to_owned(),
}),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Figure {
pub tokens: BudgetSpend,
pub cost: Option<Cost>,
pub attribution: Attribution,
}
pub fn select(bills: &[StepBill], roots: &[String]) -> Vec<StepBill> {
if roots.is_empty() {
return bills.to_vec();
}
bills
.iter()
.filter(|b| roots.iter().any(|r| Scope::Tree(r.clone()).wants(&b.conv)))
.cloned()
.collect()
}
pub fn attribution(roots: &[String]) -> Attribution {
match roots.len() {
0 => Attribution::Workspace,
n => Attribution::Conversations(n),
}
}
pub fn of_world(workspaces: &[PathBuf], prices: &Prices) -> Option<Cost> {
let mut bills = Vec::new();
for workspace in workspaces {
bills.extend(budgets::bills(workspace, &Scope::Workspace));
}
priced(&bills, prices)
}
pub fn of_conversation(bills: &[StepBill], root_id: &str, prices: &Prices) -> Figure {
let roots = [root_id.to_owned()];
figure(
&select(bills, &roots),
prices,
Attribution::Conversations(1),
)
}
pub fn of_ball(bills: &[StepBill], stamped_roots: &[String], prices: &Prices) -> Figure {
figure(
&select(bills, stamped_roots),
prices,
attribution(stamped_roots),
)
}
pub fn figure(bills: &[StepBill], prices: &Prices, attribution: Attribution) -> Figure {
Figure {
tokens: budgets::total(bills),
cost: priced(bills, prices),
attribution,
}
}
pub fn priced(bills: &[StepBill], prices: &Prices) -> Option<Cost> {
(!prices.is_empty()).then(|| cost(bills, prices))
}
fn cost(bills: &[StepBill], prices: &Prices) -> Cost {
let mut cost = Cost::default();
for bill in bills {
match prices.of(bill.model.as_deref()) {
Some(price) => cost.add(price.cost(bill.spend), 0),
None => cost.add(0, bill.spend.total_tokens()),
}
}
cost
}
#[cfg(test)]
mod tests;