diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 17f7e203877a03fc8e4e0e82245c41c96a0e4df3..061be6f738f826f3ca0c7ad8a061f7b8a402aaa6 100644 --- a/crates/agent/src/thread.rs +++ b/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, ); } diff --git a/crates/agent/src/tool_use.rs b/crates/agent/src/tool_use.rs index 32876a100c63b4cdca7d7e614e5a11ea6184e68a..2e10248e794e509bfa94d8cb74f13225bb071d03 100644 --- a/crates/agent/src/tool_use.rs +++ b/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, pending_tool_uses_by_id: HashMap, tool_result_cards: HashMap, + tool_use_metadata_by_id: HashMap, } 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, cx: &App, ) -> Option { - 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, + pub thread_id: ThreadId, + pub prompt_id: PromptId, +}