swiftide_core/
agent_traits.rs

1use std::{path::PathBuf, sync::Arc};
2
3use crate::chat_completion::ChatMessage;
4use anyhow::Result;
5use async_trait::async_trait;
6use thiserror::Error;
7
8/// A tool executor that can be used within an `AgentContext`
9#[async_trait]
10pub trait ToolExecutor: Send + Sync {
11    async fn exec_cmd(&self, cmd: &Command) -> Result<CommandOutput, CommandError>;
12}
13
14#[async_trait]
15impl<T: ToolExecutor> ToolExecutor for &T {
16    async fn exec_cmd(&self, cmd: &Command) -> Result<CommandOutput, CommandError> {
17        (*self).exec_cmd(cmd).await
18    }
19}
20
21#[async_trait]
22impl ToolExecutor for Arc<dyn ToolExecutor> {
23    async fn exec_cmd(&self, cmd: &Command) -> Result<CommandOutput, CommandError> {
24        (**self).exec_cmd(cmd).await
25    }
26}
27
28#[async_trait]
29impl ToolExecutor for Box<dyn ToolExecutor> {
30    async fn exec_cmd(&self, cmd: &Command) -> Result<CommandOutput, CommandError> {
31        (**self).exec_cmd(cmd).await
32    }
33}
34
35#[async_trait]
36impl ToolExecutor for &dyn ToolExecutor {
37    async fn exec_cmd(&self, cmd: &Command) -> Result<CommandOutput, CommandError> {
38        (**self).exec_cmd(cmd).await
39    }
40}
41
42#[derive(Debug, Error)]
43pub enum CommandError {
44    /// The executor itself failed
45    #[error("executor error: {0:#}")]
46    ExecutorError(#[from] anyhow::Error),
47
48    /// The command failed, i.e. failing tests with stderr. This error might be handled
49    #[error("command failed with NonZeroExit: {0}")]
50    NonZeroExit(CommandOutput),
51}
52
53impl From<std::io::Error> for CommandError {
54    fn from(err: std::io::Error) -> Self {
55        CommandError::NonZeroExit(err.to_string().into())
56    }
57}
58
59/// Commands that can be executed by the executor
60/// Conceptually, `Shell` allows any kind of input, and other commands enable more optimized
61/// implementations.
62///
63/// There is an ongoing consideration to make this an associated type on the executor
64///
65/// TODO: Should be able to borrow everything?
66#[non_exhaustive]
67#[derive(Debug, Clone)]
68pub enum Command {
69    Shell(String),
70    ReadFile(PathBuf),
71    WriteFile(PathBuf, String),
72}
73
74impl Command {
75    pub fn shell<S: Into<String>>(cmd: S) -> Self {
76        Command::Shell(cmd.into())
77    }
78
79    pub fn read_file<P: Into<PathBuf>>(path: P) -> Self {
80        Command::ReadFile(path.into())
81    }
82
83    pub fn write_file<P: Into<PathBuf>, S: Into<String>>(path: P, content: S) -> Self {
84        Command::WriteFile(path.into(), content.into())
85    }
86}
87
88/// Output from a `Command`
89#[derive(Debug, Clone)]
90pub struct CommandOutput {
91    pub output: String,
92    // status_code: i32,
93    // success: bool,
94}
95
96impl CommandOutput {
97    pub fn empty() -> Self {
98        CommandOutput {
99            output: String::new(),
100        }
101    }
102
103    pub fn new(output: impl Into<String>) -> Self {
104        CommandOutput {
105            output: output.into(),
106        }
107    }
108    pub fn is_empty(&self) -> bool {
109        self.output.is_empty()
110    }
111}
112
113impl std::fmt::Display for CommandOutput {
114    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
115        self.output.fmt(f)
116    }
117}
118
119impl<T: Into<String>> From<T> for CommandOutput {
120    fn from(value: T) -> Self {
121        CommandOutput {
122            output: value.into(),
123        }
124    }
125}
126
127/// Acts as the interface to the external world and manages messages for completion
128#[async_trait]
129pub trait AgentContext: Send + Sync {
130    /// List of all messages for this agent
131    ///
132    /// Used as main source for the next completion and expects all
133    /// messages to be returned if new messages are present.
134    ///
135    /// Once this method has been called, there should not be new messages
136    ///
137    /// TODO: Figure out a nice way to return a reference instead while still supporting i.e.
138    /// mutexes
139    async fn next_completion(&self) -> Option<Vec<ChatMessage>>;
140
141    /// Lists only the new messages after calling `new_completion`
142    async fn current_new_messages(&self) -> Vec<ChatMessage>;
143
144    /// Add messages for the next completion
145    async fn add_messages(&self, item: Vec<ChatMessage>);
146
147    /// Add messages for the next completion
148    async fn add_message(&self, item: ChatMessage);
149
150    /// Execute a command if the context supports it
151    async fn exec_cmd(&self, cmd: &Command) -> Result<CommandOutput, CommandError>;
152
153    async fn history(&self) -> Vec<ChatMessage>;
154
155    /// Pops the last messages up until the last completion
156    ///
157    /// LLMs failing completion for various reasons is unfortunately a common occurrence
158    /// This gives a way to redrive the last completion in a generic way
159    async fn redrive(&self);
160}
161
162#[async_trait]
163impl AgentContext for Box<dyn AgentContext> {
164    async fn next_completion(&self) -> Option<Vec<ChatMessage>> {
165        (**self).next_completion().await
166    }
167
168    async fn current_new_messages(&self) -> Vec<ChatMessage> {
169        (**self).current_new_messages().await
170    }
171
172    async fn add_messages(&self, item: Vec<ChatMessage>) {
173        (**self).add_messages(item).await;
174    }
175
176    async fn add_message(&self, item: ChatMessage) {
177        (**self).add_message(item).await;
178    }
179
180    async fn exec_cmd(&self, cmd: &Command) -> Result<CommandOutput, CommandError> {
181        (**self).exec_cmd(cmd).await
182    }
183
184    async fn history(&self) -> Vec<ChatMessage> {
185        (**self).history().await
186    }
187
188    async fn redrive(&self) {
189        (**self).redrive().await;
190    }
191}
192
193#[async_trait]
194impl AgentContext for Arc<dyn AgentContext> {
195    async fn next_completion(&self) -> Option<Vec<ChatMessage>> {
196        (**self).next_completion().await
197    }
198
199    async fn current_new_messages(&self) -> Vec<ChatMessage> {
200        (**self).current_new_messages().await
201    }
202
203    async fn add_messages(&self, item: Vec<ChatMessage>) {
204        (**self).add_messages(item).await;
205    }
206
207    async fn add_message(&self, item: ChatMessage) {
208        (**self).add_message(item).await;
209    }
210
211    async fn exec_cmd(&self, cmd: &Command) -> Result<CommandOutput, CommandError> {
212        (**self).exec_cmd(cmd).await
213    }
214
215    async fn history(&self) -> Vec<ChatMessage> {
216        (**self).history().await
217    }
218
219    async fn redrive(&self) {
220        (**self).redrive().await;
221    }
222}
223
224#[async_trait]
225impl AgentContext for &dyn AgentContext {
226    async fn next_completion(&self) -> Option<Vec<ChatMessage>> {
227        (**self).next_completion().await
228    }
229
230    async fn current_new_messages(&self) -> Vec<ChatMessage> {
231        (**self).current_new_messages().await
232    }
233
234    async fn add_messages(&self, item: Vec<ChatMessage>) {
235        (**self).add_messages(item).await;
236    }
237
238    async fn add_message(&self, item: ChatMessage) {
239        (**self).add_message(item).await;
240    }
241
242    async fn exec_cmd(&self, cmd: &Command) -> Result<CommandOutput, CommandError> {
243        (**self).exec_cmd(cmd).await
244    }
245
246    async fn history(&self) -> Vec<ChatMessage> {
247        (**self).history().await
248    }
249
250    async fn redrive(&self) {
251        (**self).redrive().await;
252    }
253}
swiftide_core/agent_traits.rs

swiftide_core/
agent_traits.rs
