From c2448e1673396d0aba47b19fbb2b2482d77f20f6 Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Tue, 3 Sep 2024 17:52:52 -0400 Subject: [PATCH] assistant: Insert creases for tool uses (#17330) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR makes it so we create creases for each of the tool uses in the context editor. Screenshot 2024-09-03 at 5 37 33 PM Release Notes: - N/A --- assets/icons/pocket_knife.svg | 1 + crates/assistant/src/assistant_panel.rs | 68 +++++++++++++++++++++++++ crates/assistant/src/context.rs | 49 +++++++++++++++++- crates/ui/src/components/icon.rs | 2 + 4 files changed, 119 insertions(+), 1 deletion(-) create mode 100644 assets/icons/pocket_knife.svg diff --git a/assets/icons/pocket_knife.svg b/assets/icons/pocket_knife.svg new file mode 100644 index 0000000000..fb2d078e20 --- /dev/null +++ b/assets/icons/pocket_knife.svg @@ -0,0 +1 @@ + diff --git a/crates/assistant/src/assistant_panel.rs b/crates/assistant/src/assistant_panel.rs index 1e95698c72..6b3a230098 100644 --- a/crates/assistant/src/assistant_panel.rs +++ b/crates/assistant/src/assistant_panel.rs @@ -1416,6 +1416,7 @@ pub struct ContextEditor { remote_id: Option, pending_slash_command_creases: HashMap, CreaseId>, pending_slash_command_blocks: HashMap, CustomBlockId>, + pending_tool_use_creases: HashMap, CreaseId>, _subscriptions: Vec, workflow_steps: HashMap, WorkflowStepViewState>, active_workflow_step: Option, @@ -1480,6 +1481,7 @@ impl ContextEditor { project, pending_slash_command_creases: HashMap::default(), pending_slash_command_blocks: HashMap::default(), + pending_tool_use_creases: HashMap::default(), _subscriptions, workflow_steps: HashMap::default(), active_workflow_step: None, @@ -1855,6 +1857,72 @@ impl ContextEditor { cx, ); } + + let new_tool_uses = self + .context + .read(cx) + .pending_tool_uses() + .into_iter() + .filter(|tool_use| { + !self + .pending_tool_use_creases + .contains_key(&tool_use.source_range) + }) + .cloned() + .collect::>(); + + let buffer = editor.buffer().read(cx).snapshot(cx); + let (excerpt_id, _buffer_id, _) = buffer.as_singleton().unwrap(); + let excerpt_id = *excerpt_id; + + let mut buffer_rows_to_fold = BTreeSet::new(); + + let creases = new_tool_uses + .iter() + .map(|tool_use| { + let placeholder = FoldPlaceholder { + render: render_fold_icon_button( + cx.view().downgrade(), + IconName::PocketKnife, + tool_use.name.clone().into(), + ), + constrain_width: false, + merge_adjacent: false, + }; + let render_trailer = + move |_row, _unfold, _cx: &mut WindowContext| Empty.into_any(); + + let start = buffer + .anchor_in_excerpt(excerpt_id, tool_use.source_range.start) + .unwrap(); + let end = buffer + .anchor_in_excerpt(excerpt_id, tool_use.source_range.end) + .unwrap(); + + let buffer_row = MultiBufferRow(start.to_point(&buffer).row); + buffer_rows_to_fold.insert(buffer_row); + + Crease::new( + start..end, + placeholder, + fold_toggle("tool-use"), + render_trailer, + ) + }) + .collect::>(); + + let crease_ids = editor.insert_creases(creases, cx); + + for buffer_row in buffer_rows_to_fold.into_iter().rev() { + editor.fold_at(&FoldAt { buffer_row }, cx); + } + + self.pending_tool_use_creases.extend( + new_tool_uses + .iter() + .map(|tool_use| tool_use.source_range.clone()) + .zip(crease_ids), + ); }); } ContextEvent::WorkflowStepsUpdated { removed, updated } => { diff --git a/crates/assistant/src/context.rs b/crates/assistant/src/context.rs index 08668c2797..cb0d45c3b3 100644 --- a/crates/assistant/src/context.rs +++ b/crates/assistant/src/context.rs @@ -490,6 +490,7 @@ pub struct Context { edits_since_last_parse: language::Subscription, finished_slash_commands: HashSet, slash_command_output_sections: Vec>, + pending_tool_uses_by_id: HashMap, message_anchors: Vec, images: HashMap, Shared>>)>, image_anchors: Vec, @@ -591,6 +592,7 @@ impl Context { messages_metadata: Default::default(), pending_slash_commands: Vec::new(), finished_slash_commands: HashSet::default(), + pending_tool_uses_by_id: HashMap::default(), slash_command_output_sections: Vec::new(), edits_since_last_parse: edits_since_last_slash_command_parse, summary: None, @@ -1004,6 +1006,14 @@ impl Context { &self.slash_command_output_sections } + pub fn pending_tool_uses(&self) -> Vec<&PendingToolUse> { + self.pending_tool_uses_by_id.values().collect() + } + + pub fn get_tool_use_by_id(&self, id: &String) -> Option<&PendingToolUse> { + self.pending_tool_uses_by_id.get(id) + } + fn set_language(&mut self, cx: &mut ModelContext) { let markdown = self.language_registry.language_for_name("Markdown"); cx.spawn(|this, mut cx| async move { @@ -1984,12 +1994,16 @@ impl Context { ); } LanguageModelCompletionEvent::ToolUse(tool_use) => { + const NEWLINE: char = '\n'; + let mut text = String::new(); - text.push('\n'); + text.push(NEWLINE); text.push_str( &serde_json::to_string_pretty(&tool_use) .expect("failed to serialize tool use to JSON"), ); + text.push(NEWLINE); + let text_len = text.len(); buffer.edit( [( @@ -1999,6 +2013,23 @@ impl Context { None, cx, ); + + let start_ix = message_old_end_offset + NEWLINE.len_utf8(); + let end_ix = + message_old_end_offset + text_len - NEWLINE.len_utf8(); + let source_range = buffer.anchor_after(start_ix) + ..buffer.anchor_after(end_ix); + + this.pending_tool_uses_by_id.insert( + tool_use.id.clone(), + PendingToolUse { + id: tool_use.id, + name: tool_use.name, + input: tool_use.input, + status: PendingToolUseStatus::Idle, + source_range, + }, + ); } } }); @@ -2757,6 +2788,22 @@ pub enum PendingSlashCommandStatus { Error(String), } +#[derive(Debug, Clone)] +pub struct PendingToolUse { + pub id: String, + pub name: String, + pub input: serde_json::Value, + pub status: PendingToolUseStatus, + pub source_range: Range, +} + +#[derive(Debug, Clone)] +pub enum PendingToolUseStatus { + Idle, + Running { _task: Shared> }, + Error(String), +} + #[derive(Serialize, Deserialize)] pub struct SavedMessage { pub id: MessageId, diff --git a/crates/ui/src/components/icon.rs b/crates/ui/src/components/icon.rs index 0d01a36e4e..0001ab4a2b 100644 --- a/crates/ui/src/components/icon.rs +++ b/crates/ui/src/components/icon.rs @@ -215,6 +215,7 @@ pub enum IconName { Pin, Play, Plus, + PocketKnife, Public, PullRequest, Quote, @@ -386,6 +387,7 @@ impl IconName { IconName::Pin => "icons/pin.svg", IconName::Play => "icons/play.svg", IconName::Plus => "icons/plus.svg", + IconName::PocketKnife => "icons/pocket_knife.svg", IconName::Public => "icons/public.svg", IconName::PullRequest => "icons/pull_request.svg", IconName::Quote => "icons/quote.svg",