omni_llm_kit/tool/
tool.rs

1use crate::common::SharedString;
2use crate::{
3    LanguageModel, LanguageModelImage, LanguageModelRequest, LanguageModelToolSchemaFormat,
4};
5use async_trait::async_trait;
6use std::fmt;
7use std::fmt::{Debug, Formatter};
8use std::ops::Deref;
9use std::pin::Pin;
10use std::sync::Arc;
11use tokio::task::JoinHandle;
12
13#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone)]
14pub enum ToolSource {
15    /// A native tool built-in to Zed.
16    Native,
17    /// A tool provided by a context server.
18    ContextServer { id: String },
19}
20/// A tool that can be used by a language model.
21pub trait Tool: Sized + Send + Sync {
22    const NAME: &'static str;
23    /// Returns the name of the tool.
24    fn name(&self) -> String {
25        Self::NAME.to_string()
26    }
27
28    /// Returns the description of the tool.
29    fn description(&self) -> String;
30
31    /// Returns the source of the tool.
32    fn source(&self) -> ToolSource {
33        ToolSource::Native
34    }
35
36    /// Returns true if the tool needs the users's confirmation
37    /// before having permission to run.
38    fn needs_confirmation(&self, input: &serde_json::Value) -> bool;
39
40    /// Returns true if the tool may perform edits.
41    fn may_perform_edits(&self) -> bool;
42
43    /// Returns the JSON schema that describes the tool's input.
44    fn input_schema(&self, _: LanguageModelToolSchemaFormat) -> anyhow::Result<serde_json::Value> {
45        Ok(serde_json::Value::Object(serde_json::Map::default()))
46    }
47    fn ui_text(&self, input: &serde_json::Value) -> String;
48    /// Runs the tool with the provided input.
49    fn run(
50        &self,
51        input: serde_json::Value,
52    ) -> impl Future<Output = anyhow::Result<ToolResultContent>> + Send;
53}
54pub trait ToolDyn: Send + Sync {
55    fn name(&self) -> String;
56    fn description(&self) -> String;
57    fn input_schema(&self, _: LanguageModelToolSchemaFormat) -> anyhow::Result<serde_json::Value>;
58    fn needs_confirmation(&self, input: &serde_json::Value) -> bool;
59    fn ui_text(&self, input: &serde_json::Value) -> String;
60    fn run(
61        &self,
62        input: serde_json::Value,
63    ) -> Pin<Box<dyn Future<Output = anyhow::Result<ToolResultContent>> + Send + '_>>;
64}
65impl<T: Tool> ToolDyn for T {
66    fn name(&self) -> String {
67        self.name()
68    }
69
70    fn description(&self) -> String {
71        self.description()
72    }
73    fn input_schema(
74        &self,
75        schema: LanguageModelToolSchemaFormat,
76    ) -> anyhow::Result<serde_json::Value> {
77        self.input_schema(schema)
78    }
79    fn ui_text(&self, input: &serde_json::Value) -> String {
80        self.ui_text(input)
81    }
82
83    fn run(
84        &self,
85        input: serde_json::Value,
86    ) -> Pin<Box<dyn Future<Output = anyhow::Result<ToolResultContent>> + Send + '_>> {
87        Box::pin(async move { <Self as Tool>::run(self, input).await })
88    }
89
90    fn needs_confirmation(&self, input: &serde_json::Value) -> bool {
91        self.needs_confirmation(input)
92    }
93}
94
95#[derive(Debug, PartialEq, Eq)]
96pub enum ToolResultContent {
97    Text(String),
98    Image(LanguageModelImage),
99}
100
101impl ToolResultContent {
102    pub fn len(&self) -> usize {
103        match self {
104            ToolResultContent::Text(str) => str.len(),
105            ToolResultContent::Image(image) => image.len(),
106        }
107    }
108
109    pub fn is_empty(&self) -> bool {
110        match self {
111            ToolResultContent::Text(str) => str.is_empty(),
112            ToolResultContent::Image(image) => image.is_empty(),
113        }
114    }
115
116    pub fn as_str(&self) -> Option<&str> {
117        match self {
118            ToolResultContent::Text(str) => Some(str),
119            ToolResultContent::Image(_) => None,
120        }
121    }
122}
123
124#[derive(Debug, Clone)]
125pub enum ToolUseStatus {
126    InputStillStreaming,
127    NeedsConfirmation,
128    Pending,
129    Running,
130    Finished(SharedString),
131    Error(SharedString),
132}
133
134impl ToolUseStatus {
135    pub fn text(&self) -> SharedString {
136        match self {
137            ToolUseStatus::NeedsConfirmation => "".into(),
138            ToolUseStatus::InputStillStreaming => "".into(),
139            ToolUseStatus::Pending => "".into(),
140            ToolUseStatus::Running => "".into(),
141            ToolUseStatus::Finished(out) => out.clone(),
142            ToolUseStatus::Error(out) => out.clone(),
143        }
144    }
145
146    pub fn error(&self) -> Option<SharedString> {
147        match self {
148            ToolUseStatus::Error(out) => Some(out.clone()),
149            _ => None,
150        }
151    }
152}
omni_llm_kit/tool/tool.rs

omni_llm_kit/tool/
tool.rs
