use crate::budgets::StepBill;
use crate::spend::Prices;
use super::StepsView;
impl StepsView {
#[must_use]
pub fn priced(mut self, bills: &[StepBill], agent: &str, prices: &Prices) -> Self {
for step in &mut self.steps {
step.cost = bills
.iter()
.find(|bill| bill.conv == agent && bill.seq == step.seq)
.and_then(|bill| crate::spend::priced(std::slice::from_ref(bill), prices));
}
self
}
}
#[cfg(test)]
mod tests {
use crate::budgets::{BudgetSpend, StepBill};
use crate::spend::Prices;
use crate::steps_view::{StepSummary, StepsView, Wound};
use serde_json::json;
fn step(seq: &str) -> StepSummary {
StepSummary {
seq: seq.to_owned(),
framing: crate::git_tree::Framing::Complete,
attempts: 1,
tokens: BudgetSpend::default(),
cost: None,
commit: None,
started_at: None,
ended_at: None,
wound: Wound::None,
}
}
fn bill(conv: &str, seq: &str, model: Option<&str>) -> StepBill {
StepBill {
conv: conv.to_owned(),
seq: seq.to_owned(),
model: model.map(str::to_owned),
provider: None,
spend: BudgetSpend {
input_tokens: 1_000_000,
..BudgetSpend::default()
},
last_usage: BudgetSpend::default(),
window: None,
wall_secs: 0,
}
}
#[test]
fn every_row_takes_its_own_bill_and_an_empty_table_prices_none() {
let view = StepsView {
steps: vec![step("001"), step("002"), step("003")],
orphan: crate::steps_view::Orphan::None,
};
let bills = [
bill("a", "001", Some("opus")),
bill("a", "002", Some("mystery")),
bill("b", "003", Some("opus")),
];
let table = Prices::from_json(&json!({ "p": { "opus": { "input": 2 } } }));
let priced = view.clone().priced(&bills, "a", &table);
let costs: Vec<Option<(u64, u64)>> = priced
.steps
.iter()
.map(|s| s.cost.map(|c| (c.micro_usd, c.unpriced_tokens)))
.collect();
assert_eq!(costs, [Some((2_000_000, 0)), Some((0, 1_000_000)), None]);
let unpriced = view.priced(&bills, "a", &Prices::default());
assert!(unpriced.steps.iter().all(|s| s.cost.is_none()));
}
}