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

Capture telemetry when requesting completions in agent2 (#36600)

Browse source 
Release Notes:

- N/A
This commit is contained in:
Antonio Scandurra 2025-08-20 18:04:32 +02:00 • committed by GitHub
parent 1e6cefaa56
commit 699f58aeba
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 26 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

1
Cargo.lock generated
View file

@ -239,6 +239,7 @@ dependencies = [
"smol",
"sqlez",
"task",
"telemetry",
"tempfile",
"terminal",
"text",

1
crates/agent2/Cargo.toml
View file

@ -58,6 +58,7 @@ settings.workspace = true
smol.workspace = true
sqlez.workspace = true
task.workspace = true
telemetry.workspace = true
terminal.workspace = true
text.workspace = true
ui.workspace = true

26
crates/agent2/src/thread.rs
View file

@ -1195,6 +1195,15 @@ impl Thread {
let mut attempt = None;
'retry: loop {
telemetry::event!(
"Agent Thread Completion",
thread_id = this.read_with(cx, |this, _| this.id.to_string())?,
prompt_id = this.read_with(cx, |this, _| this.prompt_id.to_string())?,
model = model.telemetry_id(),
model_provider = model.provider_id().to_string(),
attempt
);
let mut events = model.stream_completion(request.clone(), cx).await?;
let mut tool_uses = FuturesUnordered::new();
while let Some(event) = events.next().await {
@ -1211,8 +1220,21 @@ impl Thread {
this.update_model_request_usage(amount, limit, cx)
})?;
}
Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage)) => {
this.update(cx, |this, cx| this.update_token_usage(token_usage, cx))?;
Ok(LanguageModelCompletionEvent::UsageUpdate(usage)) => {
telemetry::event!(
"Agent Thread Completion Usage Updated",
thread_id = this.read_with(cx, |this, _| this.id.to_string())?,
prompt_id = this.read_with(cx, |this, _| this.prompt_id.to_string())?,
model = model.telemetry_id(),
model_provider = model.provider_id().to_string(),
attempt,
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,
);
this.update(cx, |this, cx| this.update_token_usage(usage, cx))?;
}
Ok(LanguageModelCompletionEvent::Stop(StopReason::Refusal)) => {
*refusal = true;