e2b863116d
Now the edit tool can access files outside the current project (just like the terminal tool can), but it's behind a prompt (unlike other edit tool actions). Release Notes: - The edit tool can now access files outside the current project, but only if the user grants it permission to.
581 lines
20 KiB
Rust
581 lines
20 KiB
Rust
use crate::{
|
|
thread::{MessageId, PromptId, ThreadId},
|
|
thread_store::SerializedMessage,
|
|
};
|
|
use agent_settings::CompletionMode;
|
|
use anyhow::Result;
|
|
use assistant_tool::{
|
|
AnyToolCard, Tool, ToolResultContent, ToolResultOutput, ToolUseStatus, ToolWorkingSet,
|
|
};
|
|
use collections::HashMap;
|
|
use futures::{FutureExt as _, future::Shared};
|
|
use gpui::{App, Entity, SharedString, Task, Window};
|
|
use icons::IconName;
|
|
use language_model::{
|
|
ConfiguredModel, LanguageModel, LanguageModelExt, LanguageModelRequest,
|
|
LanguageModelToolResult, LanguageModelToolResultContent, LanguageModelToolUse,
|
|
LanguageModelToolUseId, Role,
|
|
};
|
|
use project::Project;
|
|
use std::sync::Arc;
|
|
use util::truncate_lines_to_byte_limit;
|
|
|
|
#[derive(Debug)]
|
|
pub struct ToolUse {
|
|
pub id: LanguageModelToolUseId,
|
|
pub name: SharedString,
|
|
pub ui_text: SharedString,
|
|
pub status: ToolUseStatus,
|
|
pub input: serde_json::Value,
|
|
pub icon: icons::IconName,
|
|
pub needs_confirmation: bool,
|
|
}
|
|
|
|
pub struct ToolUseState {
|
|
tools: Entity<ToolWorkingSet>,
|
|
tool_uses_by_assistant_message: HashMap<MessageId, Vec<LanguageModelToolUse>>,
|
|
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 {
|
|
pub fn new(tools: Entity<ToolWorkingSet>) -> Self {
|
|
Self {
|
|
tools,
|
|
tool_uses_by_assistant_message: HashMap::default(),
|
|
tool_results: HashMap::default(),
|
|
pending_tool_uses_by_id: HashMap::default(),
|
|
tool_result_cards: HashMap::default(),
|
|
tool_use_metadata_by_id: HashMap::default(),
|
|
}
|
|
}
|
|
|
|
/// Constructs a [`ToolUseState`] from the given list of [`SerializedMessage`]s.
|
|
///
|
|
/// Accepts a function to filter the tools that should be used to populate the state.
|
|
///
|
|
/// If `window` is `None` (e.g., when in headless mode or when running evals),
|
|
/// tool cards won't be deserialized
|
|
pub fn from_serialized_messages(
|
|
tools: Entity<ToolWorkingSet>,
|
|
messages: &[SerializedMessage],
|
|
project: Entity<Project>,
|
|
window: Option<&mut Window>, // None in headless mode
|
|
cx: &mut App,
|
|
) -> Self {
|
|
let mut this = Self::new(tools);
|
|
let mut tool_names_by_id = HashMap::default();
|
|
let mut window = window;
|
|
|
|
for message in messages {
|
|
match message.role {
|
|
Role::Assistant => {
|
|
if !message.tool_uses.is_empty() {
|
|
let tool_uses = message
|
|
.tool_uses
|
|
.iter()
|
|
.map(|tool_use| LanguageModelToolUse {
|
|
id: tool_use.id.clone(),
|
|
name: tool_use.name.clone().into(),
|
|
raw_input: tool_use.input.to_string(),
|
|
input: tool_use.input.clone(),
|
|
is_input_complete: true,
|
|
})
|
|
.collect::<Vec<_>>();
|
|
|
|
tool_names_by_id.extend(
|
|
tool_uses
|
|
.iter()
|
|
.map(|tool_use| (tool_use.id.clone(), tool_use.name.clone())),
|
|
);
|
|
|
|
this.tool_uses_by_assistant_message
|
|
.insert(message.id, tool_uses);
|
|
|
|
for tool_result in &message.tool_results {
|
|
let tool_use_id = tool_result.tool_use_id.clone();
|
|
let Some(tool_use) = tool_names_by_id.get(&tool_use_id) else {
|
|
log::warn!("no tool name found for tool use: {tool_use_id:?}");
|
|
continue;
|
|
};
|
|
|
|
this.tool_results.insert(
|
|
tool_use_id.clone(),
|
|
LanguageModelToolResult {
|
|
tool_use_id: tool_use_id.clone(),
|
|
tool_name: tool_use.clone(),
|
|
is_error: tool_result.is_error,
|
|
content: tool_result.content.clone(),
|
|
output: tool_result.output.clone(),
|
|
},
|
|
);
|
|
|
|
if let Some(window) = &mut window {
|
|
if let Some(tool) = this.tools.read(cx).tool(tool_use, cx) {
|
|
if let Some(output) = tool_result.output.clone() {
|
|
if let Some(card) = tool.deserialize_card(
|
|
output,
|
|
project.clone(),
|
|
window,
|
|
cx,
|
|
) {
|
|
this.tool_result_cards.insert(tool_use_id, card);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Role::System | Role::User => {}
|
|
}
|
|
}
|
|
|
|
this
|
|
}
|
|
|
|
pub fn cancel_pending(&mut self) -> Vec<PendingToolUse> {
|
|
let mut cancelled_tool_uses = Vec::new();
|
|
self.pending_tool_uses_by_id
|
|
.retain(|tool_use_id, tool_use| {
|
|
if matches!(tool_use.status, PendingToolUseStatus::Error { .. }) {
|
|
return true;
|
|
}
|
|
|
|
let content = "Tool canceled by user".into();
|
|
self.tool_results.insert(
|
|
tool_use_id.clone(),
|
|
LanguageModelToolResult {
|
|
tool_use_id: tool_use_id.clone(),
|
|
tool_name: tool_use.name.clone(),
|
|
content,
|
|
output: None,
|
|
is_error: true,
|
|
},
|
|
);
|
|
cancelled_tool_uses.push(tool_use.clone());
|
|
false
|
|
});
|
|
cancelled_tool_uses
|
|
}
|
|
|
|
pub fn pending_tool_uses(&self) -> Vec<&PendingToolUse> {
|
|
self.pending_tool_uses_by_id.values().collect()
|
|
}
|
|
|
|
pub fn tool_uses_for_message(
|
|
&self,
|
|
id: MessageId,
|
|
project: &Entity<Project>,
|
|
cx: &App,
|
|
) -> Vec<ToolUse> {
|
|
let Some(tool_uses_for_message) = &self.tool_uses_by_assistant_message.get(&id) else {
|
|
return Vec::new();
|
|
};
|
|
|
|
let mut tool_uses = Vec::new();
|
|
|
|
for tool_use in tool_uses_for_message.iter() {
|
|
let tool_result = self.tool_results.get(&tool_use.id);
|
|
|
|
let status = (|| {
|
|
if let Some(tool_result) = tool_result {
|
|
let content = tool_result
|
|
.content
|
|
.to_str()
|
|
.map(|str| str.to_owned().into())
|
|
.unwrap_or_default();
|
|
|
|
return if tool_result.is_error {
|
|
ToolUseStatus::Error(content)
|
|
} else {
|
|
ToolUseStatus::Finished(content)
|
|
};
|
|
}
|
|
|
|
if let Some(pending_tool_use) = self.pending_tool_uses_by_id.get(&tool_use.id) {
|
|
match pending_tool_use.status {
|
|
PendingToolUseStatus::Idle => ToolUseStatus::Pending,
|
|
PendingToolUseStatus::NeedsConfirmation { .. } => {
|
|
ToolUseStatus::NeedsConfirmation
|
|
}
|
|
PendingToolUseStatus::Running { .. } => ToolUseStatus::Running,
|
|
PendingToolUseStatus::Error(ref err) => {
|
|
ToolUseStatus::Error(err.clone().into())
|
|
}
|
|
PendingToolUseStatus::InputStillStreaming => {
|
|
ToolUseStatus::InputStillStreaming
|
|
}
|
|
}
|
|
} else {
|
|
ToolUseStatus::Pending
|
|
}
|
|
})();
|
|
|
|
let (icon, needs_confirmation) =
|
|
if let Some(tool) = self.tools.read(cx).tool(&tool_use.name, cx) {
|
|
(
|
|
tool.icon(),
|
|
tool.needs_confirmation(&tool_use.input, project, cx),
|
|
)
|
|
} else {
|
|
(IconName::Cog, false)
|
|
};
|
|
|
|
tool_uses.push(ToolUse {
|
|
id: tool_use.id.clone(),
|
|
name: tool_use.name.clone().into(),
|
|
ui_text: self.tool_ui_label(
|
|
&tool_use.name,
|
|
&tool_use.input,
|
|
tool_use.is_input_complete,
|
|
cx,
|
|
),
|
|
input: tool_use.input.clone(),
|
|
status,
|
|
icon,
|
|
needs_confirmation,
|
|
})
|
|
}
|
|
|
|
tool_uses
|
|
}
|
|
|
|
pub fn tool_ui_label(
|
|
&self,
|
|
tool_name: &str,
|
|
input: &serde_json::Value,
|
|
is_input_complete: bool,
|
|
cx: &App,
|
|
) -> SharedString {
|
|
if let Some(tool) = self.tools.read(cx).tool(tool_name, cx) {
|
|
if is_input_complete {
|
|
tool.ui_text(input).into()
|
|
} else {
|
|
tool.still_streaming_ui_text(input).into()
|
|
}
|
|
} else {
|
|
format!("Unknown tool {tool_name:?}").into()
|
|
}
|
|
}
|
|
|
|
pub fn tool_results_for_message(
|
|
&self,
|
|
assistant_message_id: MessageId,
|
|
) -> Vec<&LanguageModelToolResult> {
|
|
let Some(tool_uses) = self
|
|
.tool_uses_by_assistant_message
|
|
.get(&assistant_message_id)
|
|
else {
|
|
return Vec::new();
|
|
};
|
|
|
|
tool_uses
|
|
.iter()
|
|
.filter_map(|tool_use| self.tool_results.get(&tool_use.id))
|
|
.collect()
|
|
}
|
|
|
|
pub fn message_has_tool_results(&self, assistant_message_id: MessageId) -> bool {
|
|
self.tool_uses_by_assistant_message
|
|
.get(&assistant_message_id)
|
|
.map_or(false, |results| !results.is_empty())
|
|
}
|
|
|
|
pub fn tool_result(
|
|
&self,
|
|
tool_use_id: &LanguageModelToolUseId,
|
|
) -> Option<&LanguageModelToolResult> {
|
|
self.tool_results.get(tool_use_id)
|
|
}
|
|
|
|
pub fn tool_result_card(&self, tool_use_id: &LanguageModelToolUseId) -> Option<&AnyToolCard> {
|
|
self.tool_result_cards.get(tool_use_id)
|
|
}
|
|
|
|
pub fn insert_tool_result_card(
|
|
&mut self,
|
|
tool_use_id: LanguageModelToolUseId,
|
|
card: AnyToolCard,
|
|
) {
|
|
self.tool_result_cards.insert(tool_use_id, card);
|
|
}
|
|
|
|
pub fn request_tool_use(
|
|
&mut self,
|
|
assistant_message_id: MessageId,
|
|
tool_use: LanguageModelToolUse,
|
|
metadata: ToolUseMetadata,
|
|
cx: &App,
|
|
) -> Arc<str> {
|
|
let tool_uses = self
|
|
.tool_uses_by_assistant_message
|
|
.entry(assistant_message_id)
|
|
.or_default();
|
|
|
|
let mut existing_tool_use_found = false;
|
|
|
|
for existing_tool_use in tool_uses.iter_mut() {
|
|
if existing_tool_use.id == tool_use.id {
|
|
*existing_tool_use = tool_use.clone();
|
|
existing_tool_use_found = true;
|
|
}
|
|
}
|
|
|
|
if !existing_tool_use_found {
|
|
tool_uses.push(tool_use.clone());
|
|
}
|
|
|
|
let status = if tool_use.is_input_complete {
|
|
self.tool_use_metadata_by_id
|
|
.insert(tool_use.id.clone(), metadata);
|
|
|
|
PendingToolUseStatus::Idle
|
|
} else {
|
|
PendingToolUseStatus::InputStillStreaming
|
|
};
|
|
|
|
let ui_text: Arc<str> = self
|
|
.tool_ui_label(
|
|
&tool_use.name,
|
|
&tool_use.input,
|
|
tool_use.is_input_complete,
|
|
cx,
|
|
)
|
|
.into();
|
|
|
|
let may_perform_edits = self
|
|
.tools
|
|
.read(cx)
|
|
.tool(&tool_use.name, cx)
|
|
.is_some_and(|tool| tool.may_perform_edits());
|
|
|
|
self.pending_tool_uses_by_id.insert(
|
|
tool_use.id.clone(),
|
|
PendingToolUse {
|
|
assistant_message_id,
|
|
id: tool_use.id,
|
|
name: tool_use.name.clone(),
|
|
ui_text: ui_text.clone(),
|
|
input: tool_use.input,
|
|
may_perform_edits,
|
|
status,
|
|
},
|
|
);
|
|
|
|
ui_text
|
|
}
|
|
|
|
pub fn run_pending_tool(
|
|
&mut self,
|
|
tool_use_id: LanguageModelToolUseId,
|
|
ui_text: SharedString,
|
|
task: Task<()>,
|
|
) {
|
|
if let Some(tool_use) = self.pending_tool_uses_by_id.get_mut(&tool_use_id) {
|
|
tool_use.ui_text = ui_text.into();
|
|
tool_use.status = PendingToolUseStatus::Running {
|
|
_task: task.shared(),
|
|
};
|
|
}
|
|
}
|
|
|
|
pub fn confirm_tool_use(
|
|
&mut self,
|
|
tool_use_id: LanguageModelToolUseId,
|
|
ui_text: impl Into<Arc<str>>,
|
|
input: serde_json::Value,
|
|
request: Arc<LanguageModelRequest>,
|
|
tool: Arc<dyn Tool>,
|
|
) {
|
|
if let Some(tool_use) = self.pending_tool_uses_by_id.get_mut(&tool_use_id) {
|
|
let ui_text = ui_text.into();
|
|
tool_use.ui_text = ui_text.clone();
|
|
let confirmation = Confirmation {
|
|
tool_use_id,
|
|
input,
|
|
request,
|
|
tool,
|
|
ui_text,
|
|
};
|
|
tool_use.status = PendingToolUseStatus::NeedsConfirmation(Arc::new(confirmation));
|
|
}
|
|
}
|
|
|
|
pub fn insert_tool_output(
|
|
&mut self,
|
|
tool_use_id: LanguageModelToolUseId,
|
|
tool_name: Arc<str>,
|
|
output: Result<ToolResultOutput>,
|
|
configured_model: Option<&ConfiguredModel>,
|
|
completion_mode: CompletionMode,
|
|
) -> Option<PendingToolUse> {
|
|
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(output) => {
|
|
let tool_result = output.content;
|
|
const BYTES_PER_TOKEN_ESTIMATE: usize = 3;
|
|
|
|
let old_use = self.pending_tool_uses_by_id.remove(&tool_use_id);
|
|
|
|
// Protect from overly large output
|
|
let tool_output_limit = configured_model
|
|
.map(|model| {
|
|
model.model.max_token_count_for_mode(completion_mode.into()) as usize
|
|
* BYTES_PER_TOKEN_ESTIMATE
|
|
})
|
|
.unwrap_or(usize::MAX);
|
|
|
|
let content = match tool_result {
|
|
ToolResultContent::Text(text) => {
|
|
let text = if text.len() < tool_output_limit {
|
|
text
|
|
} else {
|
|
let truncated = truncate_lines_to_byte_limit(&text, tool_output_limit);
|
|
format!(
|
|
"Tool result too long. The first {} bytes:\n\n{}",
|
|
truncated.len(),
|
|
truncated
|
|
)
|
|
};
|
|
LanguageModelToolResultContent::Text(text.into())
|
|
}
|
|
ToolResultContent::Image(language_model_image) => {
|
|
if language_model_image.estimate_tokens() < tool_output_limit {
|
|
LanguageModelToolResultContent::Image(language_model_image)
|
|
} else {
|
|
self.tool_results.insert(
|
|
tool_use_id.clone(),
|
|
LanguageModelToolResult {
|
|
tool_use_id: tool_use_id.clone(),
|
|
tool_name,
|
|
content: "Tool responded with an image that would exceeded the remaining tokens".into(),
|
|
is_error: true,
|
|
output: None,
|
|
},
|
|
);
|
|
|
|
return old_use;
|
|
}
|
|
}
|
|
};
|
|
|
|
self.tool_results.insert(
|
|
tool_use_id.clone(),
|
|
LanguageModelToolResult {
|
|
tool_use_id: tool_use_id.clone(),
|
|
tool_name,
|
|
content,
|
|
is_error: false,
|
|
output: output.output,
|
|
},
|
|
);
|
|
|
|
old_use
|
|
}
|
|
Err(err) => {
|
|
self.tool_results.insert(
|
|
tool_use_id.clone(),
|
|
LanguageModelToolResult {
|
|
tool_use_id: tool_use_id.clone(),
|
|
tool_name,
|
|
content: LanguageModelToolResultContent::Text(err.to_string().into()),
|
|
is_error: true,
|
|
output: None,
|
|
},
|
|
);
|
|
|
|
if let Some(tool_use) = self.pending_tool_uses_by_id.get_mut(&tool_use_id) {
|
|
tool_use.status = PendingToolUseStatus::Error(err.to_string().into());
|
|
}
|
|
|
|
self.pending_tool_uses_by_id.get(&tool_use_id).cloned()
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn has_tool_results(&self, assistant_message_id: MessageId) -> bool {
|
|
self.tool_uses_by_assistant_message
|
|
.contains_key(&assistant_message_id)
|
|
}
|
|
|
|
pub fn tool_results(
|
|
&self,
|
|
assistant_message_id: MessageId,
|
|
) -> impl Iterator<Item = (&LanguageModelToolUse, Option<&LanguageModelToolResult>)> {
|
|
self.tool_uses_by_assistant_message
|
|
.get(&assistant_message_id)
|
|
.into_iter()
|
|
.flatten()
|
|
.map(|tool_use| (tool_use, self.tool_results.get(&tool_use.id)))
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct PendingToolUse {
|
|
pub id: LanguageModelToolUseId,
|
|
/// The ID of the Assistant message in which the tool use was requested.
|
|
#[allow(unused)]
|
|
pub assistant_message_id: MessageId,
|
|
pub name: Arc<str>,
|
|
pub ui_text: Arc<str>,
|
|
pub input: serde_json::Value,
|
|
pub status: PendingToolUseStatus,
|
|
pub may_perform_edits: bool,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct Confirmation {
|
|
pub tool_use_id: LanguageModelToolUseId,
|
|
pub input: serde_json::Value,
|
|
pub ui_text: Arc<str>,
|
|
pub request: Arc<LanguageModelRequest>,
|
|
pub tool: Arc<dyn Tool>,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub enum PendingToolUseStatus {
|
|
InputStillStreaming,
|
|
Idle,
|
|
NeedsConfirmation(Arc<Confirmation>),
|
|
Running { _task: Shared<Task<()>> },
|
|
Error(#[allow(unused)] Arc<str>),
|
|
}
|
|
|
|
impl PendingToolUseStatus {
|
|
pub fn is_idle(&self) -> bool {
|
|
matches!(self, PendingToolUseStatus::Idle)
|
|
}
|
|
|
|
pub fn is_error(&self) -> bool {
|
|
matches!(self, PendingToolUseStatus::Error(_))
|
|
}
|
|
|
|
pub fn needs_confirmation(&self) -> bool {
|
|
matches!(self, PendingToolUseStatus::NeedsConfirmation { .. })
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct ToolUseMetadata {
|
|
pub model: Arc<dyn LanguageModel>,
|
|
pub thread_id: ThreadId,
|
|
pub prompt_id: PromptId,
|
|
}
|