swiftide-tasks 0.34.0

//! Adapters that turn closures and prompt-like types into task nodes.
//!
//! # Examples
//!
//! ```no_run
//! use swiftide_tasks::{NodeError, Task};
//!
//! let mut task = Task::<i32, i32>::new();
//! let start = task.register_node_fn(|input: &i32| -> Result<i32, NodeError> {
//!     Ok(*input + 1)
//! });
//!
//! task.starts_with(start);
//! ```
use std::{pin::Pin, sync::Arc};

use async_trait::async_trait;
use swiftide_core::{
    ChatCompletion, Command, CommandError, CommandOutput, SimplePrompt, ToolExecutor,
    chat_completion::{ChatCompletionRequest, ChatCompletionResponse, errors::LanguageModelError},
    prompt::Prompt,
};

use super::{
    errors::NodeError,
    node::NodeId,
    traits::{NodeArg, TaskNode},
};

/// Wraps a synchronous closure and exposes it as a [`TaskNode`].
///
/// This is the most compact way to define small task nodes in tests, examples, and simple flows.
pub struct SyncFn<F, I, O, E = NodeError>
where
    F: Fn(&I) -> Result<O, E> + Send + Sync + Clone + 'static,
    E: std::error::Error + Send + Sync + 'static,
{
    f: F,
    _phantom: std::marker::PhantomData<(I, O, E)>,
}

/// Wraps an asynchronous closure and exposes it as a [`TaskNode`].
pub struct AsyncFn<F, I, O, E = NodeError>
where
    F: for<'a> Fn(&'a I) -> Pin<Box<dyn Future<Output = Result<O, E>> + Send + 'a>>
        + Send
        + Sync
        + Clone
        + 'static,
    E: std::error::Error + Send + Sync + 'static,
{
    f: F,
    _phantom: std::marker::PhantomData<(I, O, E)>,
}

impl<F, I, O, E> Clone for SyncFn<F, I, O, E>
where
    F: Fn(&I) -> Result<O, E> + Send + Sync + Clone + 'static,
    E: std::error::Error + Send + Sync + 'static,
{
    fn clone(&self) -> Self {
        Self {
            f: self.f.clone(),
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<F, I, O, E> Clone for AsyncFn<F, I, O, E>
where
    F: for<'a> Fn(&'a I) -> Pin<Box<dyn Future<Output = Result<O, E>> + Send + 'a>>
        + Send
        + Sync
        + Clone
        + 'static,
    E: std::error::Error + Send + Sync + 'static,
{
    fn clone(&self) -> Self {
        Self {
            f: self.f.clone(),
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<F, I, O, E> SyncFn<F, I, O, E>
where
    F: Fn(&I) -> Result<O, E> + Send + Sync + Clone + 'static,
    I: NodeArg + Clone,
    O: NodeArg + Clone,
    E: std::error::Error + Send + Sync + 'static,
{
    /// Creates a new synchronous task-node adapter from `f`.
    pub fn new(f: F) -> Self {
        Self {
            f,
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<F, I, O, E> AsyncFn<F, I, O, E>
where
    F: for<'a> Fn(&'a I) -> Pin<Box<dyn Future<Output = Result<O, E>> + Send + 'a>>
        + Send
        + Sync
        + Clone
        + 'static,
    I: NodeArg + Clone,
    O: NodeArg + Clone,
    E: std::error::Error + Send + Sync + 'static,
{
    /// Creates a new asynchronous task-node adapter from `f`.
    pub fn new(f: F) -> Self {
        Self {
            f,
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<F> From<F> for SyncFn<F, (), ()>
where
    F: Fn(&()) -> Result<(), NodeError> + Send + Sync + Clone + 'static,
{
    fn from(f: F) -> Self {
        Self::new(f)
    }
}

impl<F> From<F> for AsyncFn<F, (), ()>
where
    F: for<'a> Fn(&'a ()) -> Pin<Box<dyn Future<Output = Result<(), NodeError>> + Send + 'a>>
        + Send
        + Sync
        + Clone
        + 'static,
{
    fn from(f: F) -> Self {
        Self::new(f)
    }
}

#[async_trait]
impl<F, I, O, E> TaskNode for SyncFn<F, I, O, E>
where
    F: Fn(&I) -> Result<O, E> + Clone + Send + Sync + 'static,
    I: NodeArg + Clone,
    O: NodeArg + Clone,
    E: std::error::Error + Send + Sync + 'static,
{
    type Input = I;
    type Output = O;
    type Error = E;

    async fn evaluate(
        &self,
        _node_id: &NodeId<
            dyn TaskNode<Input = Self::Input, Output = Self::Output, Error = Self::Error>,
        >,
        input: &Self::Input,
    ) -> Result<Self::Output, Self::Error> {
        (self.f)(input)
    }
}

#[async_trait]
impl<F, I, O, E> TaskNode for AsyncFn<F, I, O, E>
where
    F: for<'a> Fn(&'a I) -> Pin<Box<dyn Future<Output = Result<O, E>> + Send + 'a>>
        + Clone
        + Send
        + Sync
        + 'static,
    I: NodeArg + Clone,
    O: NodeArg + Clone,
    E: std::error::Error + Send + Sync + 'static,
{
    type Input = I;
    type Output = O;
    type Error = E;

    async fn evaluate(
        &self,
        _node_id: &NodeId<
            dyn TaskNode<Input = Self::Input, Output = Self::Output, Error = Self::Error>,
        >,
        input: &Self::Input,
    ) -> Result<Self::Output, Self::Error> {
        (self.f)(input).await
    }
}

macro_rules! impl_task_node_for_prompt_like {
    ($ty:ty, $input:ty, $output:ty, $error:ty, |$this:ident, $input_value:ident| $body:expr) => {
        #[async_trait]
        impl TaskNode for $ty {
            type Input = $input;
            type Output = $output;
            type Error = $error;

            async fn evaluate(
                &self,
                _node_id: &NodeId<
                    dyn TaskNode<Input = Self::Input, Output = Self::Output, Error = Self::Error>,
                >,
                $input_value: &Self::Input,
            ) -> Result<Self::Output, Self::Error> {
                let $this = self;
                $body
            }
        }
    };
}

impl_task_node_for_prompt_like!(
    Box<dyn SimplePrompt>,
    Prompt,
    String,
    LanguageModelError,
    |this, input| {
        // TODO: Prompt should be borrowed
        this.prompt(input.clone()).await
    }
);

impl_task_node_for_prompt_like!(
    Arc<dyn SimplePrompt>,
    Prompt,
    String,
    LanguageModelError,
    |this, input| {
        // TODO: Prompt should be borrowed
        this.prompt(input.clone()).await
    }
);

impl_task_node_for_prompt_like!(
    Box<dyn ChatCompletion>,
    ChatCompletionRequest<'static>,
    ChatCompletionResponse,
    LanguageModelError,
    |this, input| this.complete(input).await
);

impl_task_node_for_prompt_like!(
    Arc<dyn ChatCompletion>,
    ChatCompletionRequest<'static>,
    ChatCompletionResponse,
    LanguageModelError,
    |this, input| this.complete(input).await
);

impl_task_node_for_prompt_like!(
    Box<dyn ToolExecutor>,
    Command,
    CommandOutput,
    CommandError,
    |this, input| this.exec_cmd(input).await
);

impl_task_node_for_prompt_like!(
    Arc<dyn ToolExecutor>,
    Command,
    CommandOutput,
    CommandError,
    |this, input| this.exec_cmd(input).await
);