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

Add token usage to LanguageModelTextStream (#27490)

Browse source 
Release Notes:

- N/A

---------

Co-authored-by: Michael Sloan <michael@zed.dev>
This commit is contained in:
Thomas Mickley-Doyle 2025-03-26 17:21:01 -05:00 • committed by GitHub
parent 5f8c53ffe8
commit 1e8b50f471
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 49 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

38
crates/language_model/src/language_model.rs
View file

@ -14,6 +14,7 @@ use futures::FutureExt;
use futures::{future::BoxFuture, stream::BoxStream, StreamExt, TryStreamExt as _};
use gpui::{AnyElement, AnyView, App, AsyncApp, SharedString, Task, Window};
use icons::IconName;
use parking_lot::Mutex;
use proto::Plan;
use schemars::JsonSchema;
use serde::{de::DeserializeOwned, Deserialize, Serialize};
@ -141,6 +142,8 @@ pub struct LanguageModelToolUse {
pub struct LanguageModelTextStream {
pub message_id: Option<String>,
pub stream: BoxStream<'static, Result<String>>,
// Has complete token usage after the stream has finished
pub last_token_usage: Arc<Mutex<TokenUsage>>,
}
impl Default for LanguageModelTextStream {
@ -148,6 +151,7 @@ impl Default for LanguageModelTextStream {
Self {
message_id: None,
stream: Box::pin(futures::stream::empty()),
last_token_usage: Arc::new(Mutex::new(TokenUsage::default())),
}
}
}
@ -200,6 +204,7 @@ pub trait LanguageModel: Send + Sync {
let mut events = events.await?.fuse();
let mut message_id = None;
let mut first_item_text = None;
let last_token_usage = Arc::new(Mutex::new(TokenUsage::default()));
if let Some(first_event) = events.next().await {
match first_event {
@ -214,20 +219,33 @@ pub trait LanguageModel: Send + Sync {
}
let stream = futures::stream::iter(first_item_text.map(Ok))
.chain(events.filter_map(|result| async move {
match result {
Ok(LanguageModelCompletionEvent::StartMessage { .. }) => None,
Ok(LanguageModelCompletionEvent::Text(text)) => Some(Ok(text)),
Ok(LanguageModelCompletionEvent::Thinking(_)) => None,
Ok(LanguageModelCompletionEvent::Stop(_)) => None,
Ok(LanguageModelCompletionEvent::ToolUse(_)) => None,
Ok(LanguageModelCompletionEvent::UsageUpdate(_)) => None,
Err(err) => Some(Err(err)),
.chain(events.filter_map({
let last_token_usage = last_token_usage.clone();
move |result| {
let last_token_usage = last_token_usage.clone();
async move {
match result {
Ok(LanguageModelCompletionEvent::StartMessage { .. }) => None,
Ok(LanguageModelCompletionEvent::Text(text)) => Some(Ok(text)),
Ok(LanguageModelCompletionEvent::Thinking(_)) => None,
Ok(LanguageModelCompletionEvent::Stop(_)) => None,
Ok(LanguageModelCompletionEvent::ToolUse(_)) => None,
Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage)) => {
*last_token_usage.lock() = token_usage;
None
}
Err(err) => Some(Err(err)),
}
}
}
}))
.boxed();
Ok(LanguageModelTextStream { message_id, stream })
Ok(LanguageModelTextStream {
message_id,
stream,
last_token_usage,
})
}
.boxed()
}