From b25ec65b4848582344ff24fb7aaef13232ecea87 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Mon, 7 Apr 2025 21:44:05 -0400 Subject: [PATCH] Add custom render method to Tool trait --- crates/agent/src/active_thread.rs | 99 +++++++++++++++------ crates/assistant_tool/src/assistant_tool.rs | 22 ++++- 2 files changed, 95 insertions(+), 26 deletions(-) diff --git a/crates/agent/src/active_thread.rs b/crates/agent/src/active_thread.rs index 08e8e3e8da..fde445f764 100644 --- a/crates/agent/src/active_thread.rs +++ b/crates/agent/src/active_thread.rs @@ -76,6 +76,16 @@ struct RenderedToolUse { output: Entity, } +impl From<&RenderedToolUse> for assistant_tool::RenderedToolUse { + fn from(source: &RenderedToolUse) -> Self { + Self { + label: source.label.entity_id(), + input: source.input.entity_id(), + output: source.output.entity_id(), + } + } +} + impl RenderedMessage { fn from_segments( segments: &[MessageSegment], @@ -2007,24 +2017,49 @@ impl ActiveThread { results_content_container() .border_t_1() .border_color(self.tool_card_border_color(cx)) - .child( - Label::new("Result") - .size(LabelSize::XSmall) - .color(Color::Muted) - .buffer_font(cx), - ) .child(div().w_full().text_ui_sm(cx).children( rendered_tool_use.as_ref().map(|rendered| { - MarkdownElement::new( - rendered.output.clone(), - tool_use_markdown_style(window, cx), - ) - .on_url_click({ - let workspace = self.workspace.clone(); - move |text, window, cx| { - open_markdown_link(text, workspace.clone(), window, cx); + let tool_name = tool_use.name.to_string(); + let tool_registry = assistant_tool::ToolRegistry::global(cx); + + if let Some(tool) = tool_registry.tool(&tool_name) { + match tool.render(&rendered.into(), window, cx) { + Some(rendered) => rendered, + None => { + // Default to rendering the output as markdown + div() + .child( + Label::new("Result") + .size(LabelSize::XSmall) + .color(Color::Muted) + .buffer_font(cx), + ) + .child( + div().w_full().text_ui_sm(cx).child( + MarkdownElement::new( + rendered.output.clone(), + tool_use_markdown_style(window, cx), + ) + .on_url_click({ + let workspace = self.workspace.clone(); + move |text, window, cx| { + open_markdown_link( + text, + workspace.clone(), + window, + cx, + ); + } + }), + ), + ) + .into_any_element() + } } - }) + } else { + log::error!("Tool not found: {tool_name}"); + gpui::Empty.into_any_element() + } }), )), ), @@ -2071,16 +2106,30 @@ impl ActiveThread { div() .text_ui_sm(cx) .children(rendered_tool_use.as_ref().map(|rendered| { - MarkdownElement::new( - rendered.output.clone(), - tool_use_markdown_style(window, cx), - ) - .on_url_click({ - let workspace = self.workspace.clone(); - move |text, window, cx| { - open_markdown_link(text, workspace.clone(), window, cx); - } - }) + let tool_name = tool_use.name.to_string(); + let tool_registry = assistant_tool::ToolRegistry::global(cx); + + tool_registry + .tool(&tool_name) + .and_then(|tool| tool.render(&rendered.into(), window, cx)) + .unwrap_or_else(|| { + MarkdownElement::new( + rendered.output.clone(), + tool_use_markdown_style(window, cx), + ) + .on_url_click({ + let workspace = self.workspace.clone(); + move |text, window, cx| { + open_markdown_link( + text, + workspace.clone(), + window, + cx, + ); + } + }) + .into_any_element() + }) })), ), ), diff --git a/crates/assistant_tool/src/assistant_tool.rs b/crates/assistant_tool/src/assistant_tool.rs index 65c844e554..39adf84500 100644 --- a/crates/assistant_tool/src/assistant_tool.rs +++ b/crates/assistant_tool/src/assistant_tool.rs @@ -8,7 +8,7 @@ use std::fmt::Formatter; use std::sync::Arc; use anyhow::Result; -use gpui::{App, Entity, SharedString, Task}; +use gpui::{self, App, Entity, EntityId, SharedString, Task}; use icons::IconName; use language_model::LanguageModelRequestMessage; use language_model::LanguageModelToolSchemaFormat; @@ -18,6 +18,14 @@ pub use crate::action_log::*; pub use crate::tool_registry::*; pub use crate::tool_working_set::*; +/// A rendered tool use containing styled markdown elements for UI representation. +#[derive(Clone)] +pub struct RenderedToolUse { + pub label: EntityId, + pub input: EntityId, + pub output: EntityId, +} + pub fn init(cx: &mut App) { ToolRegistry::default_global(cx); } @@ -58,6 +66,18 @@ pub trait Tool: 'static + Send + Sync { /// Returns markdown to be displayed in the UI for this tool. fn ui_text(&self, input: &serde_json::Value) -> String; + /// Returns a custom UI element to render the tool's output. + /// Returns None by default to indicate that rendering has not yet been + /// implemented for this tool, and the caller should do some default rendering. + fn render( + &self, + _rendered_tool_use: &RenderedToolUse, + _window: &mut gpui::Window, + _cx: &gpui::App, + ) -> Option { + None + } + /// Runs the tool with the provided input. fn run( self: Arc,