rei-forks/fred
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

assistant2: Show token count in the Prompt Editor view (#25091)

Browse Source 
Release Notes:

- N/A

---------

Co-authored-by: Agus Zubiaga <hi@aguz.me>
This commit is contained in:
Danilo Leal 2025-02-18 13:24:15 -03:00 committed by GitHub
parent 8a7565e04b
commit 1e255e41cc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 91 additions and 61 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
crates/assistant2/src/assistant_panel.rs
View File

@ -3,8 +3,8 @@ use std::sync::Arc;
use anyhow::{anyhow, Result};
use assistant_context_editor::{
make_lsp_adapter_delegate, AssistantPanelDelegate, ConfigurationError, ContextEditor,
ContextHistory, SlashCommandCompletionProvider,
make_lsp_adapter_delegate, render_remaining_tokens, AssistantPanelDelegate, ConfigurationError,
ContextEditor, ContextHistory, SlashCommandCompletionProvider,
};
use assistant_settings::{AssistantDockPosition, AssistantSettings};
use assistant_slash_command::SlashCommandWorkingSet;
@ -635,20 +635,33 @@ impl AssistantPanel {
.border_color(cx.theme().colors().border)
.child(
h_flex()
.child(Label::new(title))
.when(sub_title.is_some(), |this| {
this.child(
h_flex()
.pl_1p5()
.gap_1p5()
.child(
Label::new("/")
.size(LabelSize::Small)
.color(Color::Disabled)
.alpha(0.5),
.w_full()
.gap_1()
.justify_between()
.child(
h_flex()
.child(Label::new(title))
.when(sub_title.is_some(), |this| {
this.child(
h_flex()
.pl_1p5()
.gap_1p5()
.child(
Label::new("/")
.size(LabelSize::Small)
.color(Color::Disabled)
.alpha(0.5),
)
.child(Label::new(sub_title.unwrap())),
)
.child(Label::new(sub_title.unwrap())),
)
}),
)
.children(if matches!(self.active_view, ActiveView::PromptEditor) {
self.context_editor
.as_ref()
.and_then(|editor| render_remaining_tokens(editor, cx))
} else {
None
}),
)
.child(

9
crates/assistant_context_editor/src/context.rs
View File

@ -1190,11 +1190,14 @@ impl AssistantContext {
let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
return;
};
let debounce = self.token_count.is_some();
self.pending_token_count = cx.spawn(|this, mut cx| {
async move {
cx.background_executor()
.timer(Duration::from_millis(200))
.await;
if debounce {
cx.background_executor()
.timer(Duration::from_millis(200))
.await;
}
let token_count = cx.update(|cx| model.count_tokens(request, cx))?.await?;
this.update(&mut cx, |this, cx| {

100
crates/assistant_context_editor/src/context_editor.rs
View File

@ -2867,7 +2867,6 @@ impl EventEmitter<SearchEvent> for ContextEditor {}
impl Render for ContextEditor {
fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
let provider = LanguageModelRegistry::read_global(cx).active_provider();
let accept_terms = if self.show_accept_terms {
provider.as_ref().and_then(|provider| {
provider.render_accept_terms(LanguageModelProviderTosView::PromptEditorPopup, cx)
@ -3250,48 +3249,58 @@ impl ContextEditorToolbarItem {
model_summary_editor,
}
}
}
fn render_remaining_tokens(&self, cx: &mut Context<Self>) -> Option<impl IntoElement> {
let context = &self
.active_context_editor
.as_ref()?
.upgrade()?
.read(cx)
.context;
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()
.child(
Label::new(humanize_token_count(token_count))
.size(LabelSize::Small)
.color(token_count_color),
)
.child(Label::new("/").size(LabelSize::Small).color(Color::Muted))
.child(
Label::new(humanize_token_count(max_token_count))
.size(LabelSize::Small)
.color(Color::Muted),
),
)
}
pub fn render_remaining_tokens(
context_editor: &Entity<ContextEditor>,
cx: &App,
) -> Option<impl IntoElement> {
let context = &context_editor.read(cx).context;
let (token_count_color, token_count, max_token_count, tooltip) = match token_state(context, cx)?
{
TokenState::NoTokensLeft {
max_token_count,
token_count,
} => (
Color::Error,
token_count,
max_token_count,
Some("Token Limit Reached"),
),
TokenState::HasMoreTokens {
max_token_count,
token_count,
over_warn_threshold,
} => {
let (color, tooltip) = if over_warn_threshold {
(Color::Warning, Some("Token Limit is Close to Exhaustion"))
} else {
(Color::Muted, None)
};
(color, token_count, max_token_count, tooltip)
}
};
Some(
h_flex()
.id("token-count")
.gap_0p5()
.child(
Label::new(humanize_token_count(token_count))
.size(LabelSize::Small)
.color(token_count_color),
)
.child(Label::new("/").size(LabelSize::Small).color(Color::Muted))
.child(
Label::new(humanize_token_count(max_token_count))
.size(LabelSize::Small)
.color(Color::Muted),
)
.when_some(tooltip, |element, tooltip| {
element.tooltip(Tooltip::text(tooltip))
}),
)
}
impl Render for ContextEditorToolbarItem {
@ -3334,7 +3343,12 @@ impl Render for ContextEditorToolbarItem {
// scan_items_remaining
// .map(|remaining_items| format!("Files to scan: {}", remaining_items))
// })
.children(self.render_remaining_tokens(cx));
.children(
self.active_context_editor
.as_ref()
.and_then(|editor| editor.upgrade())
.and_then(|editor| render_remaining_tokens(&editor, cx)),
);
h_flex()
.px_0p5()