assistant2: Automatically respond to the model with tool results (#25706)

Marshall Bowers created 10 months ago

This PR updates the tool use flow in Assistant 2 to automatically
respond to the model with tool results when the tools have finished
running.

Release Notes:

- N/A

Change summary

crates/assistant2/src/active_thread.rs | 33 +++++++++++++++++++++++++--
crates/assistant2/src/thread.rs        | 12 ++++++++++
2 files changed, 42 insertions(+), 3 deletions(-)

Detailed changes

crates/assistant2/src/active_thread.rs

@@ -8,14 +8,16 @@ use gpui::{
     UnderlineStyle, WeakEntity,
 };
 use language::LanguageRegistry;
-use language_model::{LanguageModelToolUseId, Role};
+use language_model::{LanguageModelRegistry, LanguageModelToolUseId, Role};
 use markdown::{Markdown, MarkdownStyle};
 use settings::Settings as _;
 use theme::ThemeSettings;
 use ui::{prelude::*, Disclosure};
 use workspace::Workspace;
 
-use crate::thread::{MessageId, Thread, ThreadError, ThreadEvent, ToolUse, ToolUseStatus};
+use crate::thread::{
+    MessageId, RequestKind, Thread, ThreadError, ThreadEvent, ToolUse, ToolUseStatus,
+};
 use crate::thread_store::ThreadStore;
 use crate::ui::ContextPill;
 
@@ -263,7 +265,24 @@ impl ActiveThread {
                     }
                 }
             }
-            ThreadEvent::ToolFinished { .. } => {}
+            ThreadEvent::ToolFinished { .. } => {
+                let all_tools_finished = self
+                    .thread
+                    .read(cx)
+                    .pending_tool_uses()
+                    .into_iter()
+                    .all(|tool_use| tool_use.status.is_error());
+                if all_tools_finished {
+                    let model_registry = LanguageModelRegistry::read_global(cx);
+                    if let Some(model) = model_registry.active_model() {
+                        self.thread.update(cx, |thread, cx| {
+                            // Insert an empty user message to contain the tool results.
+                            thread.insert_user_message("", Vec::new(), cx);
+                            thread.send_to_model(model, RequestKind::Chat, true, cx);
+                        });
+                    }
+                }
+            }
         }
     }
 
@@ -281,6 +300,14 @@ impl ActiveThread {
         let tool_uses = self.thread.read(cx).tool_uses_for_message(message_id);
         let colors = cx.theme().colors();
 
+        // Don't render user messages that are just there for returning tool results.
+        if message.role == Role::User
+            && message.text.is_empty()
+            && self.thread.read(cx).message_has_tool_results(message_id)
+        {
+            return Empty.into_any();
+        }
+
         let message_content = v_flex()
             .child(div().p_2p5().text_ui(cx).child(markdown.clone()))
             .when_some(context, |parent, context| {

crates/assistant2/src/thread.rs

@@ -263,6 +263,12 @@ impl Thread {
         tool_uses
     }
 
+    pub fn message_has_tool_results(&self, message_id: MessageId) -> bool {
+        self.tool_results_by_message
+            .get(&message_id)
+            .map_or(false, |results| !results.is_empty())
+    }
+
     pub fn insert_user_message(
         &mut self,
         text: impl Into<String>,
@@ -649,6 +655,8 @@ impl Thread {
                                 {
                                     tool_use.status = PendingToolUseStatus::Error(err.to_string());
                                 }
+
+                                cx.emit(ThreadEvent::ToolFinished { tool_use_id });
                             }
                         }
                     })
@@ -724,4 +732,8 @@ impl PendingToolUseStatus {
     pub fn is_idle(&self) -> bool {
         matches!(self, PendingToolUseStatus::Idle)
     }
+
+    pub fn is_error(&self) -> bool {
+        matches!(self, PendingToolUseStatus::Error(_))
+    }
 }