agent: Report usage from thread summarization requests (#29012)
This PR makes it so the thread summarization also reports the model request usage, to prevent the case where the count would appear to jump by 2 the next time a message was sent after summarization. Release Notes: - N/A
This commit is contained in:
parent
ba7f886c62
commit
676cc109a3
2 changed files with 30 additions and 9 deletions
|
@ -1302,8 +1302,15 @@ impl Thread {
|
||||||
|
|
||||||
self.pending_summary = cx.spawn(async move |this, cx| {
|
self.pending_summary = cx.spawn(async move |this, cx| {
|
||||||
async move {
|
async move {
|
||||||
let stream = model.model.stream_completion_text(request, &cx);
|
let stream = model.model.stream_completion_text_with_usage(request, &cx);
|
||||||
let mut messages = stream.await?;
|
let (mut messages, usage) = stream.await?;
|
||||||
|
|
||||||
|
if let Some(usage) = usage {
|
||||||
|
this.update(cx, |_thread, cx| {
|
||||||
|
cx.emit(ThreadEvent::UsageUpdated(usage));
|
||||||
|
})
|
||||||
|
.ok();
|
||||||
|
}
|
||||||
|
|
||||||
let mut new_summary = String::new();
|
let mut new_summary = String::new();
|
||||||
while let Some(message) = messages.stream.next().await {
|
while let Some(message) = messages.stream.next().await {
|
||||||
|
|
|
@ -262,10 +262,21 @@ pub trait LanguageModel: Send + Sync {
|
||||||
request: LanguageModelRequest,
|
request: LanguageModelRequest,
|
||||||
cx: &AsyncApp,
|
cx: &AsyncApp,
|
||||||
) -> BoxFuture<'static, Result<LanguageModelTextStream>> {
|
) -> BoxFuture<'static, Result<LanguageModelTextStream>> {
|
||||||
let events = self.stream_completion(request, cx);
|
self.stream_completion_text_with_usage(request, cx)
|
||||||
|
.map(|result| result.map(|(stream, _usage)| stream))
|
||||||
|
.boxed()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stream_completion_text_with_usage(
|
||||||
|
&self,
|
||||||
|
request: LanguageModelRequest,
|
||||||
|
cx: &AsyncApp,
|
||||||
|
) -> BoxFuture<'static, Result<(LanguageModelTextStream, Option<RequestUsage>)>> {
|
||||||
|
let future = self.stream_completion_with_usage(request, cx);
|
||||||
|
|
||||||
async move {
|
async move {
|
||||||
let mut events = events.await?.fuse();
|
let (events, usage) = future.await?;
|
||||||
|
let mut events = events.fuse();
|
||||||
let mut message_id = None;
|
let mut message_id = None;
|
||||||
let mut first_item_text = None;
|
let mut first_item_text = None;
|
||||||
let last_token_usage = Arc::new(Mutex::new(TokenUsage::default()));
|
let last_token_usage = Arc::new(Mutex::new(TokenUsage::default()));
|
||||||
|
@ -305,11 +316,14 @@ pub trait LanguageModel: Send + Sync {
|
||||||
}))
|
}))
|
||||||
.boxed();
|
.boxed();
|
||||||
|
|
||||||
Ok(LanguageModelTextStream {
|
Ok((
|
||||||
|
LanguageModelTextStream {
|
||||||
message_id,
|
message_id,
|
||||||
stream,
|
stream,
|
||||||
last_token_usage,
|
last_token_usage,
|
||||||
})
|
},
|
||||||
|
usage,
|
||||||
|
))
|
||||||
}
|
}
|
||||||
.boxed()
|
.boxed()
|
||||||
}
|
}
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue