swiftide_tasks/
adapters.rs

1//! Adapters that turn closures and prompt-like types into task nodes.
2//!
3//! # Examples
4//!
5//! ```no_run
6//! use swiftide_tasks::{NodeError, Task};
7//!
8//! let mut task = Task::<i32, i32>::new();
9//! let start = task.register_node_fn(|input: &i32| -> Result<i32, NodeError> {
10//!     Ok(*input + 1)
11//! });
12//!
13//! task.starts_with(start);
14//! ```
15use std::{pin::Pin, sync::Arc};
16
17use async_trait::async_trait;
18use swiftide_core::{
19    ChatCompletion, Command, CommandError, CommandOutput, SimplePrompt, ToolExecutor,
20    chat_completion::{ChatCompletionRequest, ChatCompletionResponse, errors::LanguageModelError},
21    prompt::Prompt,
22};
23
24use super::{
25    errors::NodeError,
26    node::NodeId,
27    traits::{NodeArg, TaskNode},
28};
29
30/// Wraps a synchronous closure and exposes it as a [`TaskNode`].
31///
32/// This is the most compact way to define small task nodes in tests, examples, and simple flows.
33pub struct SyncFn<F, I, O, E = NodeError>
34where
35    F: Fn(&I) -> Result<O, E> + Send + Sync + Clone + 'static,
36    E: std::error::Error + Send + Sync + 'static,
37{
38    f: F,
39    _phantom: std::marker::PhantomData<(I, O, E)>,
40}
41
42/// Wraps an asynchronous closure and exposes it as a [`TaskNode`].
43pub struct AsyncFn<F, I, O, E = NodeError>
44where
45    F: for<'a> Fn(&'a I) -> Pin<Box<dyn Future<Output = Result<O, E>> + Send + 'a>>
46        + Send
47        + Sync
48        + Clone
49        + 'static,
50    E: std::error::Error + Send + Sync + 'static,
51{
52    f: F,
53    _phantom: std::marker::PhantomData<(I, O, E)>,
54}
55
56impl<F, I, O, E> Clone for SyncFn<F, I, O, E>
57where
58    F: Fn(&I) -> Result<O, E> + Send + Sync + Clone + 'static,
59    E: std::error::Error + Send + Sync + 'static,
60{
61    fn clone(&self) -> Self {
62        Self {
63            f: self.f.clone(),
64            _phantom: std::marker::PhantomData,
65        }
66    }
67}
68
69impl<F, I, O, E> Clone for AsyncFn<F, I, O, E>
70where
71    F: for<'a> Fn(&'a I) -> Pin<Box<dyn Future<Output = Result<O, E>> + Send + 'a>>
72        + Send
73        + Sync
74        + Clone
75        + 'static,
76    E: std::error::Error + Send + Sync + 'static,
77{
78    fn clone(&self) -> Self {
79        Self {
80            f: self.f.clone(),
81            _phantom: std::marker::PhantomData,
82        }
83    }
84}
85
86impl<F, I, O, E> SyncFn<F, I, O, E>
87where
88    F: Fn(&I) -> Result<O, E> + Send + Sync + Clone + 'static,
89    I: NodeArg + Clone,
90    O: NodeArg + Clone,
91    E: std::error::Error + Send + Sync + 'static,
92{
93    /// Creates a new synchronous task-node adapter from `f`.
94    pub fn new(f: F) -> Self {
95        Self {
96            f,
97            _phantom: std::marker::PhantomData,
98        }
99    }
100}
101
102impl<F, I, O, E> AsyncFn<F, I, O, E>
103where
104    F: for<'a> Fn(&'a I) -> Pin<Box<dyn Future<Output = Result<O, E>> + Send + 'a>>
105        + Send
106        + Sync
107        + Clone
108        + 'static,
109    I: NodeArg + Clone,
110    O: NodeArg + Clone,
111    E: std::error::Error + Send + Sync + 'static,
112{
113    /// Creates a new asynchronous task-node adapter from `f`.
114    pub fn new(f: F) -> Self {
115        Self {
116            f,
117            _phantom: std::marker::PhantomData,
118        }
119    }
120}
121
122impl<F> From<F> for SyncFn<F, (), ()>
123where
124    F: Fn(&()) -> Result<(), NodeError> + Send + Sync + Clone + 'static,
125{
126    fn from(f: F) -> Self {
127        Self::new(f)
128    }
129}
130
131impl<F> From<F> for AsyncFn<F, (), ()>
132where
133    F: for<'a> Fn(&'a ()) -> Pin<Box<dyn Future<Output = Result<(), NodeError>> + Send + 'a>>
134        + Send
135        + Sync
136        + Clone
137        + 'static,
138{
139    fn from(f: F) -> Self {
140        Self::new(f)
141    }
142}
143
144#[async_trait]
145impl<F, I, O, E> TaskNode for SyncFn<F, I, O, E>
146where
147    F: Fn(&I) -> Result<O, E> + Clone + Send + Sync + 'static,
148    I: NodeArg + Clone,
149    O: NodeArg + Clone,
150    E: std::error::Error + Send + Sync + 'static,
151{
152    type Input = I;
153    type Output = O;
154    type Error = E;
155
156    async fn evaluate(
157        &self,
158        _node_id: &NodeId<
159            dyn TaskNode<Input = Self::Input, Output = Self::Output, Error = Self::Error>,
160        >,
161        input: &Self::Input,
162    ) -> Result<Self::Output, Self::Error> {
163        (self.f)(input)
164    }
165}
166
167#[async_trait]
168impl<F, I, O, E> TaskNode for AsyncFn<F, I, O, E>
169where
170    F: for<'a> Fn(&'a I) -> Pin<Box<dyn Future<Output = Result<O, E>> + Send + 'a>>
171        + Clone
172        + Send
173        + Sync
174        + 'static,
175    I: NodeArg + Clone,
176    O: NodeArg + Clone,
177    E: std::error::Error + Send + Sync + 'static,
178{
179    type Input = I;
180    type Output = O;
181    type Error = E;
182
183    async fn evaluate(
184        &self,
185        _node_id: &NodeId<
186            dyn TaskNode<Input = Self::Input, Output = Self::Output, Error = Self::Error>,
187        >,
188        input: &Self::Input,
189    ) -> Result<Self::Output, Self::Error> {
190        (self.f)(input).await
191    }
192}
193
194macro_rules! impl_task_node_for_prompt_like {
195    ($ty:ty, $input:ty, $output:ty, $error:ty, |$this:ident, $input_value:ident| $body:expr) => {
196        #[async_trait]
197        impl TaskNode for $ty {
198            type Input = $input;
199            type Output = $output;
200            type Error = $error;
201
202            async fn evaluate(
203                &self,
204                _node_id: &NodeId<
205                    dyn TaskNode<Input = Self::Input, Output = Self::Output, Error = Self::Error>,
206                >,
207                $input_value: &Self::Input,
208            ) -> Result<Self::Output, Self::Error> {
209                let $this = self;
210                $body
211            }
212        }
213    };
214}
215
216impl_task_node_for_prompt_like!(
217    Box<dyn SimplePrompt>,
218    Prompt,
219    String,
220    LanguageModelError,
221    |this, input| {
222        // TODO: Prompt should be borrowed
223        this.prompt(input.clone()).await
224    }
225);
226
227impl_task_node_for_prompt_like!(
228    Arc<dyn SimplePrompt>,
229    Prompt,
230    String,
231    LanguageModelError,
232    |this, input| {
233        // TODO: Prompt should be borrowed
234        this.prompt(input.clone()).await
235    }
236);
237
238impl_task_node_for_prompt_like!(
239    Box<dyn ChatCompletion>,
240    ChatCompletionRequest<'static>,
241    ChatCompletionResponse,
242    LanguageModelError,
243    |this, input| this.complete(input).await
244);
245
246impl_task_node_for_prompt_like!(
247    Arc<dyn ChatCompletion>,
248    ChatCompletionRequest<'static>,
249    ChatCompletionResponse,
250    LanguageModelError,
251    |this, input| this.complete(input).await
252);
253
254impl_task_node_for_prompt_like!(
255    Box<dyn ToolExecutor>,
256    Command,
257    CommandOutput,
258    CommandError,
259    |this, input| this.exec_cmd(input).await
260);
261
262impl_task_node_for_prompt_like!(
263    Arc<dyn ToolExecutor>,
264    Command,
265    CommandOutput,
266    CommandError,
267    |this, input| this.exec_cmd(input).await
268);
swiftide_tasks/adapters.rs

swiftide_tasks/
adapters.rs
