assistant_tool.rs

  1mod action_log;
  2mod tool_registry;
  3mod tool_schema;
  4mod tool_working_set;
  5
  6use std::fmt;
  7use std::fmt::Debug;
  8use std::fmt::Formatter;
  9use std::sync::Arc;
 10
 11use anyhow::Result;
 12use gpui::AnyElement;
 13use gpui::AnyWindowHandle;
 14use gpui::Context;
 15use gpui::IntoElement;
 16use gpui::Window;
 17use gpui::{App, Entity, SharedString, Task, WeakEntity};
 18use icons::IconName;
 19use language_model::LanguageModelRequestMessage;
 20use language_model::LanguageModelToolSchemaFormat;
 21use project::Project;
 22use workspace::Workspace;
 23
 24pub use crate::action_log::*;
 25pub use crate::tool_registry::*;
 26pub use crate::tool_schema::*;
 27pub use crate::tool_working_set::*;
 28
 29pub fn init(cx: &mut App) {
 30    ToolRegistry::default_global(cx);
 31}
 32
 33#[derive(Debug, Clone)]
 34pub enum ToolUseStatus {
 35    InputStillStreaming,
 36    NeedsConfirmation,
 37    Pending,
 38    Running,
 39    Finished(SharedString),
 40    Error(SharedString),
 41}
 42
 43impl ToolUseStatus {
 44    pub fn text(&self) -> SharedString {
 45        match self {
 46            ToolUseStatus::NeedsConfirmation => "".into(),
 47            ToolUseStatus::InputStillStreaming => "".into(),
 48            ToolUseStatus::Pending => "".into(),
 49            ToolUseStatus::Running => "".into(),
 50            ToolUseStatus::Finished(out) => out.clone(),
 51            ToolUseStatus::Error(out) => out.clone(),
 52        }
 53    }
 54
 55    pub fn error(&self) -> Option<SharedString> {
 56        match self {
 57            ToolUseStatus::Error(out) => Some(out.clone()),
 58            _ => None,
 59        }
 60    }
 61}
 62
 63/// The result of running a tool, containing both the asynchronous output
 64/// and an optional card view that can be rendered immediately.
 65pub struct ToolResult {
 66    /// The asynchronous task that will eventually resolve to the tool's output
 67    pub output: Task<Result<String>>,
 68    /// An optional view to present the output of the tool.
 69    pub card: Option<AnyToolCard>,
 70}
 71
 72pub trait ToolCard: 'static + Sized {
 73    fn render(
 74        &mut self,
 75        status: &ToolUseStatus,
 76        window: &mut Window,
 77        workspace: WeakEntity<Workspace>,
 78        cx: &mut Context<Self>,
 79    ) -> impl IntoElement;
 80}
 81
 82#[derive(Clone)]
 83pub struct AnyToolCard {
 84    entity: gpui::AnyEntity,
 85    render: fn(
 86        entity: gpui::AnyEntity,
 87        status: &ToolUseStatus,
 88        window: &mut Window,
 89        workspace: WeakEntity<Workspace>,
 90        cx: &mut App,
 91    ) -> AnyElement,
 92}
 93
 94impl<T: ToolCard> From<Entity<T>> for AnyToolCard {
 95    fn from(entity: Entity<T>) -> Self {
 96        fn downcast_render<T: ToolCard>(
 97            entity: gpui::AnyEntity,
 98            status: &ToolUseStatus,
 99            window: &mut Window,
100            workspace: WeakEntity<Workspace>,
101            cx: &mut App,
102        ) -> AnyElement {
103            let entity = entity.downcast::<T>().unwrap();
104            entity.update(cx, |entity, cx| {
105                entity
106                    .render(status, window, workspace, cx)
107                    .into_any_element()
108            })
109        }
110
111        Self {
112            entity: entity.into(),
113            render: downcast_render::<T>,
114        }
115    }
116}
117
118impl AnyToolCard {
119    pub fn render(
120        &self,
121        status: &ToolUseStatus,
122        window: &mut Window,
123        workspace: WeakEntity<Workspace>,
124        cx: &mut App,
125    ) -> AnyElement {
126        (self.render)(self.entity.clone(), status, window, workspace, cx)
127    }
128}
129
130impl From<Task<Result<String>>> for ToolResult {
131    /// Convert from a task to a ToolResult with no card
132    fn from(output: Task<Result<String>>) -> Self {
133        Self { output, card: None }
134    }
135}
136
137#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone)]
138pub enum ToolSource {
139    /// A native tool built-in to Zed.
140    Native,
141    /// A tool provided by a context server.
142    ContextServer { id: SharedString },
143}
144
145/// A tool that can be used by a language model.
146pub trait Tool: 'static + Send + Sync {
147    /// Returns the name of the tool.
148    fn name(&self) -> String;
149
150    /// Returns the description of the tool.
151    fn description(&self) -> String;
152
153    /// Returns the icon for the tool.
154    fn icon(&self) -> IconName;
155
156    /// Returns the source of the tool.
157    fn source(&self) -> ToolSource {
158        ToolSource::Native
159    }
160
161    /// Returns true iff the tool needs the users's confirmation
162    /// before having permission to run.
163    fn needs_confirmation(&self, input: &serde_json::Value, cx: &App) -> bool;
164
165    /// Returns the JSON schema that describes the tool's input.
166    fn input_schema(&self, _: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
167        Ok(serde_json::Value::Object(serde_json::Map::default()))
168    }
169
170    /// Returns markdown to be displayed in the UI for this tool.
171    fn ui_text(&self, input: &serde_json::Value) -> String;
172
173    /// Returns markdown to be displayed in the UI for this tool, while the input JSON is still streaming
174    /// (so information may be missing).
175    fn still_streaming_ui_text(&self, input: &serde_json::Value) -> String {
176        self.ui_text(input)
177    }
178
179    /// Runs the tool with the provided input.
180    fn run(
181        self: Arc<Self>,
182        input: serde_json::Value,
183        messages: &[LanguageModelRequestMessage],
184        project: Entity<Project>,
185        action_log: Entity<ActionLog>,
186        window: Option<AnyWindowHandle>,
187        cx: &mut App,
188    ) -> ToolResult;
189}
190
191impl Debug for dyn Tool {
192    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
193        f.debug_struct("Tool").field("name", &self.name()).finish()
194    }
195}