use std::{borrow::Cow, sync::Arc};
use async_trait::async_trait;
use derive_builder::Builder;
use serde::Deserialize;
use swiftide::{
chat_completion::{self, errors::ToolError, Tool, ToolOutput, ToolSpec},
traits::AgentContext,
};
use crate::agent::{running_agent::RunningAgent, session::Session};
#[derive(Clone, Builder)]
pub struct DelegateAgent {
session: Arc<Session>,
agent: RunningAgent,
tool_spec: ToolSpec,
}
impl DelegateAgent {
#[must_use]
pub fn builder() -> DelegateAgentBuilder {
DelegateAgentBuilder::default()
}
pub async fn delegate_agent(
&self,
_context: &dyn AgentContext,
task: &str,
) -> Result<ToolOutput, ToolError> {
self.session.swap_agent(self.agent.clone())?;
self.agent.query(task).await?;
tracing::info!("Delegated task to agent");
Ok(ToolOutput::Stop)
}
}
#[derive(Deserialize)]
struct DelegateArgs {
task: String,
}
#[async_trait]
impl Tool for DelegateAgent {
async fn invoke(
&self,
agent_context: &dyn AgentContext,
raw_args: Option<&str>,
) -> Result<ToolOutput, ToolError> {
let Some(args) = raw_args else {
return Err(ToolError::MissingArguments(format!(
"No arguments provided for {}",
self.name()
)));
};
let args: DelegateArgs = serde_json::from_str(&args)?;
return self.delegate_agent(agent_context, &args.task).await;
}
fn tool_spec(&self) -> chat_completion::ToolSpec {
self.tool_spec.clone()
}
fn name(&self) -> Cow<'_, str> {
self.tool_spec().name.into()
}
}