llimorse_tools/
subagent.rs1use 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
11pub trait ToolFactory: fmt::Debug + Send + Sync {
13 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 #[derive(Debug)]
50 'params: pub struct SubagentParams {
51 prompt: String,
53
54 #[serde(default)]
58 fork: bool,
59 }
60
61 'result: pub struct SubagentResult {
63 response: String,
65 }
66
67 #[derive(Debug)]
69 'state: pub struct Subagent<F: ToolFactory, N: SubagentNotifier> {
70 system_prompt: String,
72
73 llama_url: String,
75
76 model_name: String,
78
79 tool_factory: F,
81
82 notifier: N,
84
85 concurrency: Option<Arc<Semaphore>>,
87 }
88}
89
90impl<F: ToolFactory, N: SubagentNotifier> Subagent<F, N> {
91 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 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 agent.push_user(format!("{}\n\n{}", self.system_prompt, params.prompt));
155 } else {
156 agent.push_system(&self.system_prompt);
157 agent.push_user(¶ms.prompt);
158 }
159
160 self.tool_factory.add_tools(&mut agent);
161
162 let connector = self
163 .notifier
164 .created(&agent, ¶ms.prompt)
165 .await
166 .map_err(Into::into)?;
167
168 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}