Skip to main content

swiftide_agents/
task_agent.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use swiftide_core::prompt::Prompt;
5use swiftide_tasks::{NodeId, TaskNode};
6use tokio::sync::Mutex;
7
8use crate::{Agent, errors::AgentError};
9
10/// Wraps an [`Agent`](crate::Agent) so it can participate in a task graph.
11///
12/// For more control you can always implement [`TaskNode`] directly.
13///
14/// `TaskAgent` keeps the wrapped agent behind a mutex and therefore serializes access when the
15/// same adapter instance is shared across multiple task branches.
16#[derive(Clone, Debug)]
17pub struct TaskAgent(Arc<Mutex<Agent>>);
18
19impl From<Agent> for TaskAgent {
20    fn from(agent: Agent) -> Self {
21        Self(Arc::new(Mutex::new(agent)))
22    }
23}
24
25#[async_trait]
26impl TaskNode for TaskAgent {
27    type Input = Prompt;
28    type Output = ();
29    type Error = AgentError;
30
31    async fn evaluate(
32        &self,
33        _node_id: &NodeId<
34            dyn TaskNode<Input = Self::Input, Output = Self::Output, Error = Self::Error>,
35        >,
36        input: &Self::Input,
37    ) -> Result<Self::Output, Self::Error> {
38        self.0.lock().await.query(input.clone()).await
39    }
40}