Skip to main content

llimorse_tools/
subagent.rs

1//! Subagent tool: delegate tasks to a secondary agent.
2
3use anyhow::{Result, anyhow};
4use futures::StreamExt;
5use llimorse::line_format::{ChatMessage, ToolCall};
6use llimorse::{Agent, CallableTool, StreamingChunk};
7use std::fmt;
8use std::sync::Arc;
9use tokio::sync::Semaphore;
10
11/// Factory for creating the tool set that subagents receive.
12pub trait ToolFactory: fmt::Debug + Send + Sync {
13    /// Add a fresh set of tools to a new subagent.
14    fn add_tools(&self, agent: &mut Agent);
15}
16
17#[allow(async_fn_in_trait)]
18pub trait SubagentNotifier: fmt::Debug + Send {
19    type Error: Into<anyhow::Error> + Send + Sync + 'static;
20    type Subagent: SubagentConnector;
21
22    async fn created(
23        &self,
24        agent: &Agent,
25        prompt: &str,
26    ) -> std::result::Result<Self::Subagent, Self::Error>;
27}
28
29#[allow(async_fn_in_trait)]
30pub trait SubagentConnector {
31    type Error: Into<anyhow::Error>;
32
33    async fn push_chunk(&self, chunk: StreamingChunk);
34    fn tool_call(&self, agent: &Agent, call: &ToolCall) -> std::result::Result<(), Self::Error>;
35    fn tool_result(
36        &self,
37        agent: &Agent,
38        call: &ToolCall,
39        result: &Result<String, anyhow::Error>,
40    ) -> std::result::Result<(), Self::Error>;
41}
42
43llimorse::tool! {
44    'name: "subagent";
45
46    /// Delegate a task to a subagent to explore questions or execute self-contained tasks without
47    /// polluting your context with unrelated tool call results. The subagent will return its final
48    /// response.
49    #[derive(Debug)]
50    'params: pub struct SubagentParams {
51        /// The prompt to give to the subagent, which is all it will receive from you
52        prompt: String,
53
54        /// If true, the subagent inherits the parent agent's full conversation history, giving it
55        /// complete context. If false (default), the subagent starts with a clean context
56        /// containing only the prompt.
57        #[serde(default)]
58        fork: bool,
59    }
60
61    /// Result of running a subagent
62    'result: pub struct SubagentResult {
63        /// The subagent's final response
64        response: String,
65    }
66
67    /// Delegate tasks to subagents.
68    #[derive(Debug)]
69    'state: pub struct Subagent<F: ToolFactory, N: SubagentNotifier> {
70        /// System prompt pushed into every subagent conversation.
71        system_prompt: String,
72
73        /// llama-server base URL (e.g. "http://127.0.0.1:8080").
74        llama_url: String,
75
76        /// Model name to use (from the parent's [`ClientState`](llimorse::client::ClientState)).
77        model_name: String,
78
79        /// Factory producing the tool set for each subagent.
80        tool_factory: F,
81
82        /// Notify e.g. a UI of changes in subagents
83        notifier: N,
84
85        /// Limits the number of subagents that may run in parallel; `None` means unlimited.
86        concurrency: Option<Arc<Semaphore>>,
87    }
88}
89
90impl<F: ToolFactory, N: SubagentNotifier> Subagent<F, N> {
91    /// Create a new subagent tool.
92    pub fn new(
93        system_prompt: &str,
94        llama_url: &str,
95        model_name: &str,
96        tool_factory: F,
97        notifier: N,
98        max_parallel: Option<usize>,
99    ) -> Self {
100        Subagent {
101            system_prompt: system_prompt.to_string(),
102            llama_url: llama_url.to_string(),
103            model_name: model_name.to_string(),
104            tool_factory,
105            notifier,
106            concurrency: max_parallel.map(|limit| Arc::new(Semaphore::new(limit))),
107        }
108    }
109}
110
111impl fmt::Display for SubagentParams {
112    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
113        if self.fork {
114            write!(f, "[fork] {}", self.prompt)
115        } else {
116            write!(f, "{}", self.prompt)
117        }
118    }
119}
120
121impl fmt::Display for SubagentResult {
122    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
123        write!(f, "{}", self.response)
124    }
125}
126
127impl<F: ToolFactory, N: SubagentNotifier> CallableTool for Subagent<F, N> {
128    async fn execute(&self, parent: &Agent, params: SubagentParams) -> Result<SubagentResult> {
129        let client = llimorse::Client::new(&self.llama_url, Some(&self.model_name))
130            .await
131            .map_err(|err| {
132                anyhow!(
133                    "Failed to connect to {} to launch a subagent: {err}",
134                    self.llama_url
135                )
136            })?;
137        let mut agent = llimorse::Agent::new(client);
138
139        if params.fork {
140            // The parent's last message contains this very tool call (and possibly others); drop
141            // the calls, so the history does not end in unanswered calls, and drop the whole
142            // message if nothing remains
143            let mut history = parent.history().to_vec();
144            if let Some(ChatMessage::Assistant(msg)) = history.last_mut() {
145                msg.tool_calls = None;
146            }
147            if history.last().is_some_and(ChatMessage::is_empty) {
148                history.pop();
149            }
150            agent.push_history(history);
151
152            // Not all chat templates support system messages in the middle of the conversation,
153            // so pass the instructions as part of the user message
154            agent.push_user(format!("{}\n\n{}", self.system_prompt, params.prompt));
155        } else {
156            agent.push_system(&self.system_prompt);
157            agent.push_user(&params.prompt);
158        }
159
160        self.tool_factory.add_tools(&mut agent);
161
162        let connector = self
163            .notifier
164            .created(&agent, &params.prompt)
165            .await
166            .map_err(Into::into)?;
167
168        // The permit is held for the whole run, so at most `max_parallel` subagents run at once;
169        // it is released when the guard is dropped at the end of this function
170        // Registering with the notifier happens before taking the permit, so subagents waiting for
171        // a free slot are already visible to the UI (with their prompt), just not running.
172        let _permit = match &self.concurrency {
173            Some(semaphore) => {
174                Some(semaphore.clone().acquire_owned().await.map_err(|err| {
175                    anyhow!("Failed to acquire subagent concurrency permit: {err}")
176                })?)
177            }
178            None => None,
179        };
180
181        loop {
182            let mut stream = agent.submit().await?;
183            while let Some(chunk) = stream.next().await {
184                connector.push_chunk(chunk?).await;
185            }
186            drop(stream);
187
188            if !agent
189                .execute_pending_calls(
190                    |agent, call| connector.tool_call(agent, call).map_err(Into::into),
191                    |agent, call, result| {
192                        connector
193                            .tool_result(agent, call, result)
194                            .map_err(Into::into)
195                    },
196                )
197                .await
198            {
199                break;
200            }
201        }
202
203        Ok(SubagentResult {
204            response: agent
205                .last_result()
206                .cloned()
207                .ok_or_else(|| anyhow!("Subagent produced no response"))?,
208        })
209    }
210}