agent: Add additional fields to `Agent Tool Finished` telemetry event (#29163)

Marshall Bowers created 8 months ago

This PR adds additional fields to the `Agent Tool Finished` telemetry
event:

- `model`
- `model_provider`
- `thread_id`
- `prompt_id`

Release Notes:

- N/A

Change summary

crates/agent/src/thread.rs   |  9 ++++++++-
crates/agent/src/tool_use.rs | 33 ++++++++++++++++++++++++++++++---
2 files changed, 38 insertions(+), 4 deletions(-)

Detailed changes

crates/agent/src/thread.rs

@@ -38,7 +38,7 @@ use crate::thread_store::{
     SerializedMessage, SerializedMessageSegment, SerializedThread, SerializedToolResult,
     SerializedToolUse, SharedProjectContext,
 };
-use crate::tool_use::{PendingToolUse, ToolUse, ToolUseState, USING_TOOL_MARKER};
+use crate::tool_use::{PendingToolUse, ToolUse, ToolUseMetadata, ToolUseState, USING_TOOL_MARKER};
 
 #[derive(
     Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize, JsonSchema,
@@ -1180,6 +1180,12 @@ impl Thread {
             None
         };
         let prompt_id = self.last_prompt_id.clone();
+        let tool_use_metadata = ToolUseMetadata {
+            model: model.clone(),
+            thread_id: self.id.clone(),
+            prompt_id: prompt_id.clone(),
+        };
+
         let task = cx.spawn(async move |thread, cx| {
             let stream_completion_future = model.stream_completion_with_usage(request, &cx);
             let initial_token_usage =
@@ -1288,6 +1294,7 @@ impl Thread {
                                 thread.tool_use.request_tool_use(
                                     last_assistant_message_id,
                                     tool_use,
+                                    tool_use_metadata.clone(),
                                     cx,
                                 );
                             }

crates/agent/src/tool_use.rs

@@ -7,13 +7,13 @@ use futures::FutureExt as _;
 use futures::future::Shared;
 use gpui::{App, Entity, SharedString, Task};
 use language_model::{
-    LanguageModelRegistry, LanguageModelRequestMessage, LanguageModelToolResult,
+    LanguageModel, LanguageModelRegistry, LanguageModelRequestMessage, LanguageModelToolResult,
     LanguageModelToolUse, LanguageModelToolUseId, MessageContent, Role,
 };
 use ui::IconName;
 use util::truncate_lines_to_byte_limit;
 
-use crate::thread::MessageId;
+use crate::thread::{MessageId, PromptId, ThreadId};
 use crate::thread_store::SerializedMessage;
 
 #[derive(Debug)]
@@ -36,6 +36,7 @@ pub struct ToolUseState {
     tool_results: HashMap<LanguageModelToolUseId, LanguageModelToolResult>,
     pending_tool_uses_by_id: HashMap<LanguageModelToolUseId, PendingToolUse>,
     tool_result_cards: HashMap<LanguageModelToolUseId, AnyToolCard>,
+    tool_use_metadata_by_id: HashMap<LanguageModelToolUseId, ToolUseMetadata>,
 }
 
 impl ToolUseState {
@@ -47,6 +48,7 @@ impl ToolUseState {
             tool_results: HashMap::default(),
             pending_tool_uses_by_id: HashMap::default(),
             tool_result_cards: HashMap::default(),
+            tool_use_metadata_by_id: HashMap::default(),
         }
     }
 
@@ -254,6 +256,7 @@ impl ToolUseState {
         &mut self,
         assistant_message_id: MessageId,
         tool_use: LanguageModelToolUse,
+        metadata: ToolUseMetadata,
         cx: &App,
     ) {
         self.tool_uses_by_assistant_message
@@ -261,6 +264,9 @@ impl ToolUseState {
             .or_default()
             .push(tool_use.clone());
 
+        self.tool_use_metadata_by_id
+            .insert(tool_use.id.clone(), metadata);
+
         // The tool use is being requested by the Assistant, so we want to
         // attach the tool results to the next user message.
         let next_user_message_id = MessageId(assistant_message_id.0 + 1);
@@ -327,7 +333,21 @@ impl ToolUseState {
         output: Result<String>,
         cx: &App,
     ) -> Option<PendingToolUse> {
-        telemetry::event!("Agent Tool Finished", tool_name, success = output.is_ok());
+        let metadata = self.tool_use_metadata_by_id.remove(&tool_use_id);
+
+        telemetry::event!(
+            "Agent Tool Finished",
+            model = metadata
+                .as_ref()
+                .map(|metadata| metadata.model.telemetry_id()),
+            model_provider = metadata
+                .as_ref()
+                .map(|metadata| metadata.model.provider_id().to_string()),
+            thread_id = metadata.as_ref().map(|metadata| metadata.thread_id.clone()),
+            prompt_id = metadata.as_ref().map(|metadata| metadata.prompt_id.clone()),
+            tool_name,
+            success = output.is_ok()
+        );
 
         match output {
             Ok(tool_result) => {
@@ -496,3 +516,10 @@ impl PendingToolUseStatus {
         matches!(self, PendingToolUseStatus::NeedsConfirmation { .. })
     }
 }
+
+#[derive(Clone)]
+pub struct ToolUseMetadata {
+    pub model: Arc<dyn LanguageModel>,
+    pub thread_id: ThreadId,
+    pub prompt_id: PromptId,
+}