Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

agent: Log errors from token counting APIs (#29626)

Browse source 
Release Notes:

- N/A
This commit is contained in:
Michael Sloan 2025-04-29 15:55:54 -06:00 • committed by GitHub
parent f7a3e00074
commit 33abf1ee7c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 100 additions and 86 deletions
Download patch file Download diff file Expand all files Collapse all files

93
crates/agent/src/active_thread.rs
View file

@ -702,7 +702,7 @@ struct EditMessageState {
editor: Entity<Editor>,
last_estimated_token_count: Option<usize>,
_subscription: Subscription,
_update_token_count_task: Option<Task<anyhow::Result<()>>>,
_update_token_count_task: Option<Task<()>>,
}
impl ActiveThread {
@ -1268,58 +1268,65 @@ impl ActiveThread {
.await;
}
let token_count = if let Some(task) = cx.update(|cx| {
let Some(message) = thread.read(cx).message(message_id) else {
log::error!("Message that was being edited no longer exists");
return None;
};
let message_text = editor.read(cx).text(cx);
let token_count = if let Some(task) = cx
.update(|cx| {
let Some(message) = thread.read(cx).message(message_id) else {
log::error!("Message that was being edited no longer exists");
return None;
};
let message_text = editor.read(cx).text(cx);
if message_text.is_empty() && message.loaded_context.is_empty() {
return None;
}
if message_text.is_empty() && message.loaded_context.is_empty() {
return None;
}
let mut request_message = LanguageModelRequestMessage {
role: language_model::Role::User,
content: Vec::new(),
cache: false,
};
let mut request_message = LanguageModelRequestMessage {
role: language_model::Role::User,
content: Vec::new(),
cache: false,
};
message
.loaded_context
.add_to_request_message(&mut request_message);
message
.loaded_context
.add_to_request_message(&mut request_message);
if !message_text.is_empty() {
request_message
.content
.push(MessageContent::Text(message_text));
}
if !message_text.is_empty() {
request_message
.content
.push(MessageContent::Text(message_text));
}
let request = language_model::LanguageModelRequest {
thread_id: None,
prompt_id: None,
mode: None,
messages: vec![request_message],
tools: vec![],
stop: vec![],
temperature: None,
};
let request = language_model::LanguageModelRequest {
thread_id: None,
prompt_id: None,
mode: None,
messages: vec![request_message],
tools: vec![],
stop: vec![],
temperature: None,
};
Some(configured_model.model.count_tokens(request, cx))
})? {
task.await?
Some(configured_model.model.count_tokens(request, cx))
})
.ok()
.flatten()
{
task.await.log_err()
} else {
0
Some(0)
};
this.update(cx, |this, cx| {
let Some((_message_id, state)) = this.editing_message.as_mut() else {
return;
};
if let Some(token_count) = token_count {
this.update(cx, |this, cx| {
let Some((_message_id, state)) = this.editing_message.as_mut() else {
return;
};
state.last_estimated_token_count = Some(token_count);
cx.emit(ActiveThreadEvent::EditingMessageTokenCountChanged);
})
state.last_estimated_token_count = Some(token_count);
cx.emit(ActiveThreadEvent::EditingMessageTokenCountChanged);
})
.ok();
};
}));
}

93
crates/agent/src/message_editor.rs
View file

@ -63,7 +63,7 @@ pub struct MessageEditor {
edits_expanded: bool,
editor_is_expanded: bool,
last_estimated_token_count: Option<usize>,
update_token_count_task: Option<Task<anyhow::Result<()>>>,
update_token_count_task: Option<Task<()>>,
_subscriptions: Vec<Subscription>,
}
@ -1088,57 +1088,64 @@ impl MessageEditor {
.await;
}
let token_count = if let Some(task) = this.update(cx, |this, cx| {
let loaded_context = this
.last_loaded_context
.as_ref()
.map(|context_load_result| &context_load_result.loaded_context);
let message_text = editor.read(cx).text(cx);
let token_count = if let Some(task) = this
.update(cx, |this, cx| {
let loaded_context = this
.last_loaded_context
.as_ref()
.map(|context_load_result| &context_load_result.loaded_context);
let message_text = editor.read(cx).text(cx);
if message_text.is_empty()
&& loaded_context.map_or(true, |loaded_context| loaded_context.is_empty())
{
return None;
}
if message_text.is_empty()
&& loaded_context.map_or(true, |loaded_context| loaded_context.is_empty())
{
return None;
}
let mut request_message = LanguageModelRequestMessage {
role: language_model::Role::User,
content: Vec::new(),
cache: false,
};
let mut request_message = LanguageModelRequestMessage {
role: language_model::Role::User,
content: Vec::new(),
cache: false,
};
if let Some(loaded_context) = loaded_context {
loaded_context.add_to_request_message(&mut request_message);
}
if let Some(loaded_context) = loaded_context {
loaded_context.add_to_request_message(&mut request_message);
}
if !message_text.is_empty() {
request_message
.content
.push(MessageContent::Text(message_text));
}
if !message_text.is_empty() {
request_message
.content
.push(MessageContent::Text(message_text));
}
let request = language_model::LanguageModelRequest {
thread_id: None,
prompt_id: None,
mode: None,
messages: vec![request_message],
tools: vec![],
stop: vec![],
temperature: None,
};
let request = language_model::LanguageModelRequest {
thread_id: None,
prompt_id: None,
mode: None,
messages: vec![request_message],
tools: vec![],
stop: vec![],
temperature: None,
};
Some(model.model.count_tokens(request, cx))
})? {
task.await?
Some(model.model.count_tokens(request, cx))
})
.ok()
.flatten()
{
task.await.log_err()
} else {
0
Some(0)
};
this.update(cx, |this, cx| {
this.last_estimated_token_count = Some(token_count);
cx.emit(MessageEditorEvent::EstimatedTokenCount);
this.update_token_count_task.take();
})
if let Some(token_count) = token_count {
this.update(cx, |this, cx| {
this.last_estimated_token_count = Some(token_count);
cx.emit(MessageEditorEvent::EstimatedTokenCount);
this.update_token_count_task.take();
})
.ok();
}
}));
}
}