assistant_tool.rs

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