Capture telemetry when requesting completions in agent2 (#36600)

Antonio Scandurra created 3 months ago

Release Notes:

- N/A

Change summary

Cargo.lock                  |  1 +
crates/agent2/Cargo.toml    |  1 +
crates/agent2/src/thread.rs | 26 ++++++++++++++++++++++++--
3 files changed, 26 insertions(+), 2 deletions(-)

Detailed changes

Cargo.lock

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

crates/agent2/Cargo.toml

@@ -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

crates/agent2/src/thread.rs

@@ -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;