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

Record token usage telemetry (#26962)

Browse source 
<img width="1103" alt="Screenshot 2025-03-17 at 9 47 32 PM"
src="https://github.com/user-attachments/assets/947cf33d-4464-4305-8ff0-3630529d2f81"
/>


Release Notes:

- N/A
This commit is contained in:
Richard Feldman 2025-03-19 10:47:46 -04:00 • committed by GitHub
parent 3edf930007
commit 6303751325
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 17 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

17
crates/assistant2/src/thread.rs
View file

@ -604,6 +604,8 @@ impl Thread {
let task = cx.spawn(async move |thread, cx| {
let stream = model.stream_completion(request, &cx);
let initial_token_usage =
thread.read_with(&cx, |thread, _cx| thread.cumulative_token_usage.clone());
let stream_completion = async {
let mut events = stream.await?;
let mut stop_reason = StopReason::EndTurn;
@ -718,6 +720,21 @@ impl Thread {
}
}
cx.emit(ThreadEvent::DoneStreaming);
if let Ok(initial_usage) = initial_token_usage {
let usage = thread.cumulative_token_usage.clone() - initial_usage;
telemetry::event!(
"Assistant Thread Completion",
thread_id = thread.id().to_string(),
model = model.telemetry_id(),
model_provider = model.provider_id().to_string(),
input_tokens = usage.input_tokens,
output_tokens = usage.output_tokens,
cache_creation_input_tokens = usage.cache_creation_input_tokens,
cache_read_input_tokens = usage.cache_read_input_tokens,
);
}
})
.ok();
});