Capture telemetry when requesting completions in agent2 (#36600)
Release Notes: - N/A
This commit is contained in:
parent
1e6cefaa56
commit
699f58aeba
3 changed files with 26 additions and 2 deletions
1
Cargo.lock
generated
1
Cargo.lock
generated
|
@ -239,6 +239,7 @@ dependencies = [
|
||||||
"smol",
|
"smol",
|
||||||
"sqlez",
|
"sqlez",
|
||||||
"task",
|
"task",
|
||||||
|
"telemetry",
|
||||||
"tempfile",
|
"tempfile",
|
||||||
"terminal",
|
"terminal",
|
||||||
"text",
|
"text",
|
||||||
|
|
|
@ -58,6 +58,7 @@ settings.workspace = true
|
||||||
smol.workspace = true
|
smol.workspace = true
|
||||||
sqlez.workspace = true
|
sqlez.workspace = true
|
||||||
task.workspace = true
|
task.workspace = true
|
||||||
|
telemetry.workspace = true
|
||||||
terminal.workspace = true
|
terminal.workspace = true
|
||||||
text.workspace = true
|
text.workspace = true
|
||||||
ui.workspace = true
|
ui.workspace = true
|
||||||
|
|
|
@ -1195,6 +1195,15 @@ impl Thread {
|
||||||
|
|
||||||
let mut attempt = None;
|
let mut attempt = None;
|
||||||
'retry: loop {
|
'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 events = model.stream_completion(request.clone(), cx).await?;
|
||||||
let mut tool_uses = FuturesUnordered::new();
|
let mut tool_uses = FuturesUnordered::new();
|
||||||
while let Some(event) = events.next().await {
|
while let Some(event) = events.next().await {
|
||||||
|
@ -1211,8 +1220,21 @@ impl Thread {
|
||||||
this.update_model_request_usage(amount, limit, cx)
|
this.update_model_request_usage(amount, limit, cx)
|
||||||
})?;
|
})?;
|
||||||
}
|
}
|
||||||
Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage)) => {
|
Ok(LanguageModelCompletionEvent::UsageUpdate(usage)) => {
|
||||||
this.update(cx, |this, cx| this.update_token_usage(token_usage, cx))?;
|
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)) => {
|
Ok(LanguageModelCompletionEvent::Stop(StopReason::Refusal)) => {
|
||||||
*refusal = true;
|
*refusal = true;
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue