diff --git a/crates/assistant/src/assistant_panel.rs b/crates/assistant/src/assistant_panel.rs index 5d9e2e84a3..ae39254980 100644 --- a/crates/assistant/src/assistant_panel.rs +++ b/crates/assistant/src/assistant_panel.rs @@ -58,6 +58,7 @@ use std::{ time::Duration, }; use terminal_view::{terminal_panel::TerminalPanel, TerminalView}; +use ui::TintColor; use ui::{ prelude::*, utils::{format_distance_from_now, DateTimeType}, @@ -2320,8 +2321,34 @@ impl ContextEditor { }, None => "Send", }; + + let (style, tooltip) = match token_state(&self.context, cx) { + Some(TokenState::NoTokensLeft { .. }) => ( + ButtonStyle::Tinted(TintColor::Negative), + Some(Tooltip::text("Token limit reached", cx)), + ), + Some(TokenState::HasMoreTokens { + over_warn_threshold, + .. + }) => { + let (style, tooltip) = if over_warn_threshold { + ( + ButtonStyle::Tinted(TintColor::Warning), + Some(Tooltip::text("Token limit is close to exhaustion", cx)), + ) + } else { + (ButtonStyle::Filled, None) + }; + (style, tooltip) + } + None => (ButtonStyle::Filled, None), + }; + ButtonLike::new("send_button") - .style(ButtonStyle::Filled) + .style(style) + .when_some(tooltip, |button, tooltip| { + button.tooltip(move |_| tooltip.clone()) + }) .layer(ElevationIndex::ModalSurface) .children( KeyBinding::for_action_in(&Assist, &focus_handle, cx) @@ -2726,25 +2753,30 @@ impl ContextEditorToolbarItem { } fn render_remaining_tokens(&self, cx: &mut ViewContext) -> Option { - let model = LanguageModelRegistry::read_global(cx).active_model()?; let context = &self .active_context_editor .as_ref()? .upgrade()? .read(cx) .context; - let token_count = context.read(cx).token_count()?; - let max_token_count = model.max_token_count(); - - let remaining_tokens = max_token_count as isize - token_count as isize; - let token_count_color = if remaining_tokens <= 0 { - Color::Error - } else if token_count as f32 / max_token_count as f32 >= 0.8 { - Color::Warning - } else { - Color::Muted + let (token_count_color, token_count, max_token_count) = match token_state(context, cx)? { + TokenState::NoTokensLeft { + max_token_count, + token_count, + } => (Color::Error, token_count, max_token_count), + TokenState::HasMoreTokens { + max_token_count, + token_count, + over_warn_threshold, + } => { + let color = if over_warn_threshold { + Color::Warning + } else { + Color::Muted + }; + (color, token_count, max_token_count) + } }; - Some( h_flex() .gap_0p5() @@ -3077,3 +3109,40 @@ fn slash_command_error_block_renderer(message: String) -> RenderBlock { .into_any() }) } + +enum TokenState { + NoTokensLeft { + max_token_count: usize, + token_count: usize, + }, + HasMoreTokens { + max_token_count: usize, + token_count: usize, + over_warn_threshold: bool, + }, +} + +fn token_state(context: &Model, cx: &AppContext) -> Option { + const WARNING_TOKEN_THRESHOLD: f32 = 0.8; + + let model = LanguageModelRegistry::read_global(cx).active_model()?; + let token_count = context.read(cx).token_count()?; + let max_token_count = model.max_token_count(); + + let remaining_tokens = max_token_count as isize - token_count as isize; + let token_state = if remaining_tokens <= 0 { + TokenState::NoTokensLeft { + max_token_count, + token_count, + } + } else { + let over_warn_threshold = + token_count as f32 / max_token_count as f32 >= WARNING_TOKEN_THRESHOLD; + TokenState::HasMoreTokens { + max_token_count, + token_count, + over_warn_threshold, + } + }; + Some(token_state) +}