guise-ui 1.5.2

//! `AITokenMeter` — how much of the context window is gone.
//!
//! Context exhaustion is the failure people hit without warning: the
//! conversation works, and then it doesn't. A bar that turns amber at
//! three-quarters and red at ninety percent turns that into something visible
//! a few turns ahead of time, which is the only point at which it's still
//! cheap to do something about it.

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

use crate::devtools::Probed;
use crate::feedback::Progress;
use crate::theme::{theme, Color, Size};

/// Fraction of the window at which the meter warns.
const WARN_AT: f32 = 0.75;
/// Fraction at which it goes critical.
const DANGER_AT: f32 = 0.90;

/// A used/limit gauge for a model's context window.
#[derive(IntoElement)]
pub struct AITokenMeter {
  used: u64,
  limit: u64,
  label: Option<SharedString>,
  size: Size,
  /// Draw the bar as well as the numbers.
  bar: bool,
}

impl AITokenMeter {
  pub fn new(used: u64, limit: u64) -> Self {
    AITokenMeter {
      used,
      limit,
      label: None,
      size: Size::Xs,
      bar: true,
    }
  }

  /// Text before the count — "Context", "Session".
  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
  }

  /// Numbers only, for a status bar with no room for a bar.
  pub fn bar(mut self, bar: bool) -> Self {
    self.bar = bar;
    self
  }

  /// How full the window is, in `0.0..=1.0`. A zero limit reads as empty
  /// rather than dividing by it.
  pub fn fraction(&self) -> f32 {
    if self.limit == 0 {
      return 0.0;
    }
    (self.used as f64 / self.limit as f64).clamp(0.0, 1.0) as f32
  }
}

impl RenderOnce for AITokenMeter {
  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 fraction = self.fraction();

    let fill: Color = if fraction >= DANGER_AT {
      t.danger()
    } else if fraction >= WARN_AT {
      t.warning()
    } else {
      t.primary()
    };
    let text_color = if fraction >= WARN_AT {
      fill.hsla()
    } else {
      dimmed
    };

    div()
      .flex()
      .flex_col()
      .gap(px(3.0))
      .w_full()
      .child(
        div()
          .flex()
          .items_center()
          .justify_between()
          .gap(px(8.0))
          .text_size(px(font))
          .text_color(dimmed)
          .children(self.label)
          .child(
            div()
              .text_color(text_color)
              .child(SharedString::from(format!(
                "{} / {}",
                compact(self.used),
                compact(self.limit)
              ))),
          ),
      )
      .when(self.bar, |column| {
        column.child(
          Progress::new(fraction * 100.0)
            .size(Size::Xs)
            .color(fill.hsla()),
        )
      })
      .probe("AITokenMeter")
  }
}

/// Token counts run to six digits, which is unreadable in a status bar.
pub(crate) fn compact(n: u64) -> String {
  match n {
    0..=999 => n.to_string(),
    1_000..=999_999 => {
      let thousands = n as f64 / 1_000.0;
      if thousands < 10.0 {
        format!("{thousands:.1}k")
      } else {
        format!("{}k", thousands.round() as u64)
      }
    }
    _ => {
      let millions = n as f64 / 1_000_000.0;
      format!("{millions:.1}M")
    }
  }
}

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

  #[test]
  fn compact_counts_stay_short() {
    assert_eq!(compact(0), "0");
    assert_eq!(compact(999), "999");
    assert_eq!(compact(1_000), "1.0k");
    assert_eq!(compact(9_949), "9.9k");
    assert_eq!(compact(12_500), "13k");
    assert_eq!(compact(200_000), "200k");
    assert_eq!(compact(1_500_000), "1.5M");
  }

  #[test]
  fn a_zero_limit_reads_as_empty_rather_than_dividing_by_it() {
    assert_eq!(AITokenMeter::new(10, 0).fraction(), 0.0);
    assert_eq!(AITokenMeter::new(0, 100).fraction(), 0.0);
  }

  #[test]
  fn overflowing_the_window_pins_at_full() {
    assert_eq!(AITokenMeter::new(300, 100).fraction(), 1.0);
    assert_eq!(AITokenMeter::new(50, 100).fraction(), 0.5);
  }

  #[test]
  fn huge_counts_do_not_overflow_the_fraction() {
    assert_eq!(AITokenMeter::new(u64::MAX, u64::MAX).fraction(), 1.0);
    assert!(AITokenMeter::new(1, u64::MAX).fraction() < 0.001);
  }
}