guise-ui 1.5.0

//! `AICost` — what this conversation has cost so far.
//!
//! Per-token prices are small enough to feel free and add up fast enough to
//! surprise, so a running total is worth the corner it takes. The arithmetic
//! is here rather than in the caller because getting it wrong by a factor of a
//! thousand is easy: prices are quoted per million tokens.

use gpui::prelude::*;
use gpui::{div, px, App, IntoElement, SharedString, Window};

use super::tokenmeter::compact;
use crate::devtools::Probed;
use crate::icon::{Icon, IconName};
use crate::theme::{theme, ColorName, Size};

/// Prices in US dollars per million tokens, the unit providers quote.
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct AIPricing {
    pub input_per_million: f64,
    pub output_per_million: f64,
    /// What a cache read costs, when the provider bills it separately.
    pub cache_read_per_million: f64,
}

impl AIPricing {
    pub fn new(input_per_million: f64, output_per_million: f64) -> Self {
        AIPricing {
            input_per_million,
            output_per_million,
            cache_read_per_million: 0.0,
        }
    }

    pub fn cache_read(mut self, per_million: f64) -> Self {
        self.cache_read_per_million = per_million;
        self
    }
}

/// A token tally for one request or a whole session.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AIUsage {
    pub input: u64,
    pub output: u64,
    pub cache_read: u64,
}

impl AIUsage {
    pub fn new(input: u64, output: u64) -> Self {
        AIUsage {
            input,
            output,
            cache_read: 0,
        }
    }

    pub fn cache_read(mut self, tokens: u64) -> Self {
        self.cache_read = tokens;
        self
    }

    /// Every token that passed through, for a context-window meter.
    pub fn total(self) -> u64 {
        self.input
            .saturating_add(self.output)
            .saturating_add(self.cache_read)
    }

    /// What this usage costs in dollars at these prices.
    pub fn cost(self, pricing: AIPricing) -> f64 {
        let per = |tokens: u64, price: f64| tokens as f64 / 1_000_000.0 * price;
        per(self.input, pricing.input_per_million)
            + per(self.output, pricing.output_per_million)
            + per(self.cache_read, pricing.cache_read_per_million)
    }
}

/// Session totals are the sum of every request's usage. Saturating rather
/// than wrapping: a tally that silently rolls over to nearly zero is worse
/// than one that sticks at the ceiling.
impl std::ops::Add for AIUsage {
    type Output = AIUsage;

    fn add(self, other: AIUsage) -> AIUsage {
        AIUsage {
            input: self.input.saturating_add(other.input),
            output: self.output.saturating_add(other.output),
            cache_read: self.cache_read.saturating_add(other.cache_read),
        }
    }
}

impl std::iter::Sum for AIUsage {
    fn sum<I: Iterator<Item = AIUsage>>(iter: I) -> AIUsage {
        iter.fold(AIUsage::default(), |total, usage| total + usage)
    }
}

/// A running cost readout.
#[derive(IntoElement)]
pub struct AICost {
    usage: AIUsage,
    pricing: AIPricing,
    label: Option<SharedString>,
    size: Size,
    /// Show the input/output split under the total.
    breakdown: bool,
}

impl AICost {
    pub fn new(usage: AIUsage, pricing: AIPricing) -> Self {
        AICost {
            usage,
            pricing,
            label: None,
            size: Size::Xs,
            breakdown: false,
        }
    }

    pub fn label(mut self, label: impl Into<SharedString>) -> Self {
        self.label = Some(label.into());
        self
    }

    pub fn size(mut self, size: Size) -> Self {
        self.size = size;
        self
    }

    /// Break the total down by token kind.
    pub fn breakdown(mut self, breakdown: bool) -> Self {
        self.breakdown = breakdown;
        self
    }

    /// The total in dollars.
    pub fn total(&self) -> f64 {
        self.usage.cost(self.pricing)
    }
}

/// Money, at a precision that suits the amount: fractions of a cent still
/// need to be legible, dollars don't need six decimal places.
pub fn format_cost(dollars: f64) -> String {
    if !dollars.is_finite() || dollars <= 0.0 {
        return "$0.00".to_string();
    }
    if dollars < 0.01 {
        format!("${dollars:.4}")
    } else if dollars < 1.0 {
        format!("${dollars:.3}")
    } else {
        format!("${dollars:.2}")
    }
}

impl RenderOnce for AICost {
    fn render(self, _window: &mut Window, cx: &mut App) -> impl IntoElement {
        let t = theme(cx);
        let font = t.font_size(self.size);
        let dimmed = t.dimmed().hsla();
        let text_color = t.text().hsla();
        let total = self.total();

        div()
            .flex()
            .flex_col()
            .gap(px(2.0))
            .text_size(px(font))
            .child(
                div()
                    .flex()
                    .items_center()
                    .gap(px(6.0))
                    .text_color(dimmed)
                    .child(
                        Icon::new(IconName::Coins)
                            .size(Size::Xs)
                            .color(ColorName::Gray),
                    )
                    .children(self.label)
                    .child(
                        div()
                            .text_color(text_color)
                            .child(SharedString::from(format_cost(total))),
                    ),
            )
            .when(self.breakdown, |column| {
                column.child(
                    div()
                        .flex()
                        .gap(px(10.0))
                        .text_color(dimmed)
                        .child(SharedString::from(format!(
                            "in {}",
                            compact(self.usage.input)
                        )))
                        .child(SharedString::from(format!(
                            "out {}",
                            compact(self.usage.output)
                        )))
                        .when(self.usage.cache_read > 0, |row| {
                            row.child(SharedString::from(format!(
                                "cache {}",
                                compact(self.usage.cache_read)
                            )))
                        }),
                )
            })
            .probe("AICost")
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn cost_is_per_million_not_per_token() {
        // A million input tokens at $3/M is exactly $3.
        let usage = AIUsage::new(1_000_000, 0);
        assert_eq!(usage.cost(AIPricing::new(3.0, 15.0)), 3.0);
        let usage = AIUsage::new(0, 1_000_000);
        assert_eq!(usage.cost(AIPricing::new(3.0, 15.0)), 15.0);
    }

    #[test]
    fn cache_reads_bill_at_their_own_rate() {
        let pricing = AIPricing::new(3.0, 15.0).cache_read(0.3);
        let usage = AIUsage::new(0, 0).cache_read(2_000_000);
        assert!((usage.cost(pricing) - 0.6).abs() < 1e-9);
    }

    #[test]
    fn usage_adds_without_overflowing() {
        let huge = AIUsage {
            input: u64::MAX,
            output: u64::MAX,
            cache_read: u64::MAX,
        };
        assert_eq!((huge + huge).input, u64::MAX);
        assert_eq!(huge.total(), u64::MAX);
        // Summing a session's requests uses the same saturating add.
        let total: AIUsage = [AIUsage::new(1, 2), AIUsage::new(3, 4)].into_iter().sum();
        assert_eq!(total, AIUsage::new(4, 6));
    }

    #[test]
    fn tiny_amounts_stay_legible_and_bad_input_reads_as_zero() {
        assert_eq!(format_cost(0.0), "$0.00");
        assert_eq!(format_cost(-1.0), "$0.00");
        assert_eq!(format_cost(f64::NAN), "$0.00");
        assert_eq!(format_cost(0.00042), "$0.0004");
        assert_eq!(format_cost(0.125), "$0.125");
        assert_eq!(format_cost(12.3456), "$12.35");
    }
}