tiny_loop/agent.rs
1use crate::{
2 history::{History, InfiniteHistory},
3 llm::LLMProvider,
4 tool::{ClosureTool, ParallelExecutor, ToolArgs, ToolExecutor},
5 types::ToolDefinition,
6};
7
8/// Agent loop that coordinates LLM calls and tool execution.
9/// Uses [`ParallelExecutor`] by default.
10pub struct Agent {
11 pub history: Box<dyn History>,
12 llm: Box<dyn LLMProvider>,
13 executor: Box<dyn ToolExecutor>,
14 tools: Vec<ToolDefinition>,
15}
16
17impl Agent {
18 /// Create a new agent loop
19 pub fn new(llm: impl LLMProvider + 'static) -> Self {
20 Self {
21 llm: Box::new(llm),
22 history: Box::new(InfiniteHistory::new()),
23 executor: Box::new(ParallelExecutor::new()),
24 tools: Vec::new(),
25 }
26 }
27
28 /// Set custom history manager (default: [`InfiniteHistory`])
29 ///
30 /// # Example
31 /// ```
32 /// use tiny_loop::{Agent, history::InfiniteHistory, llm::OpenAIProvider};
33 ///
34 /// let agent = Agent::new(OpenAIProvider::new())
35 /// .history(InfiniteHistory::new());
36 /// ```
37 pub fn history(mut self, history: impl History + 'static) -> Self {
38 self.history = Box::new(history);
39 self
40 }
41
42 /// Append a system message
43 ///
44 /// # Example
45 /// ```
46 /// use tiny_loop::{Agent, llm::OpenAIProvider};
47 ///
48 /// let agent = Agent::new(OpenAIProvider::new())
49 /// .system("You are a helpful assistant");
50 /// ```
51 pub fn system(mut self, content: impl Into<String>) -> Self {
52 self.history.add(crate::types::TimedMessage {
53 message: crate::types::SystemMessage {
54 content: content.into(),
55 }
56 .into(),
57 timestamp: std::time::SystemTime::now(),
58 elapsed: std::time::Duration::ZERO,
59 });
60 self
61 }
62
63 /// Get reference to registered tool definitions
64 pub fn tools(&self) -> &[ToolDefinition] {
65 &self.tools
66 }
67
68 /// Set a custom tool executor (default: [`ParallelExecutor`])
69 ///
70 /// # Example
71 /// ```
72 /// use tiny_loop::{Agent, tool::SequentialExecutor, llm::OpenAIProvider};
73 ///
74 /// let agent = Agent::new(OpenAIProvider::new())
75 /// .executor(SequentialExecutor::new());
76 /// ```
77 pub fn executor(mut self, executor: impl ToolExecutor + 'static) -> Self {
78 self.executor = Box::new(executor);
79 self
80 }
81
82 /// Register a tool function created by [`#[tool]`](crate::tool::tool)
83 ///
84 /// To register a tool method with an instance, use [`Self::bind`].
85 /// To register external tools (e.g. from MCP servers) use [`Self::external`]
86 ///
87 /// # Example
88 /// ```
89 /// use tiny_loop::{Agent, tool::tool, llm::OpenAIProvider};
90 ///
91 /// #[tool]
92 /// async fn fetch(
93 /// /// URL to fetch
94 /// url: String,
95 /// ) -> String {
96 /// todo!()
97 /// }
98 ///
99 /// let agent = Agent::new(OpenAIProvider::new())
100 /// .tool(fetch);
101 /// ```
102 pub fn tool<Args, Fut>(mut self, tool: fn(Args) -> Fut) -> Self
103 where
104 Fut: Future<Output = String> + Send + 'static,
105 Args: ToolArgs + 'static,
106 {
107 self.tools.push(Args::definition());
108 self.executor.add(
109 Args::TOOL_NAME.into(),
110 Box::new(ClosureTool::boxed(move |s: String| {
111 Box::pin(async move {
112 let args = match serde_json::from_str::<Args>(&s) {
113 Ok(args) => args,
114 Err(e) => return e.to_string(),
115 };
116 tool(args).await
117 })
118 })),
119 );
120 self
121 }
122
123 /// Bind an instance to a tool method created by [`#[tool]`](crate::tool::tool)
124 ///
125 /// To register a standalone tool function, use [`Self::tool`].
126 /// To register external tools (e.g. from MCP servers) use [`Self::external`]
127 ///
128 /// # Example
129 /// ```
130 /// use tiny_loop::{Agent, tool::tool, llm::OpenAIProvider};
131 /// use std::sync::Arc;
132 ///
133 /// #[derive(Clone)]
134 /// struct Database {
135 /// data: Arc<String>,
136 /// }
137 ///
138 /// #[tool]
139 /// impl Database {
140 /// /// Fetch data from database
141 /// async fn fetch(
142 /// self,
143 /// /// Data key
144 /// key: String,
145 /// ) -> String {
146 /// todo!()
147 /// }
148 /// }
149 ///
150 /// let db = Database { data: Arc::new("data".into()) };
151 /// let agent = Agent::new(OpenAIProvider::new())
152 /// .bind(db, Database::fetch);
153 /// ```
154 pub fn bind<T, Args, Fut>(mut self, ins: T, tool: fn(T, Args) -> Fut) -> Self
155 where
156 T: Send + Sync + Clone + 'static,
157 Fut: Future<Output = String> + Send + 'static,
158 Args: ToolArgs + 'static,
159 {
160 self.tools.push(Args::definition());
161 self.executor.add(
162 Args::TOOL_NAME.into(),
163 Box::new(ClosureTool::boxed(move |s: String| {
164 let ins = ins.clone();
165 Box::pin(async move {
166 let args = match serde_json::from_str::<Args>(&s) {
167 Ok(args) => args,
168 Err(e) => return e.to_string(),
169 };
170 tool(ins, args).await
171 })
172 })),
173 );
174 self
175 }
176
177 /// Register external tools (e.g. from MCP servers)
178 ///
179 /// To register a standalone tool function, use [`tool`](Self::tool).
180 /// To register a tool method with an instance, use [`bind`](Self::bind).
181 ///
182 /// # Example
183 /// ```
184 /// use tiny_loop::{Agent, llm::OpenAIProvider, types::{Parameters, ToolDefinition, ToolFunction}};
185 /// use serde_json::{json, Value};
186 ///
187 /// let defs = vec![ToolDefinition {
188 /// tool_type: "function".into(),
189 /// function: ToolFunction {
190 /// name: "get_weather".into(),
191 /// description: "Get weather information".into(),
192 /// parameters: Parameters::from_object(
193 /// json!({
194 /// "type": "object",
195 /// "properties": {
196 /// "city": {
197 /// "type": "string",
198 /// "description": "City name"
199 /// }
200 /// },
201 /// "required": ["city"]
202 /// }).as_object().unwrap().clone()
203 /// ),
204 /// },
205 /// }];
206 ///
207 /// let external_executor = move |name: String, args: String| {
208 /// async move {
209 /// let _args = serde_json::from_str::<Value>(&args).unwrap();
210 /// "result".into()
211 /// }
212 /// };
213 ///
214 /// let agent = Agent::new(OpenAIProvider::new())
215 /// .external(defs, external_executor);
216 /// ```
217 pub fn external<Fut>(
218 mut self,
219 defs: Vec<ToolDefinition>,
220 exec: impl Fn(String, String) -> Fut + Clone + Send + Sync + 'static,
221 ) -> Self
222 where
223 Fut: Future<Output = String> + Send + 'static,
224 {
225 for d in &defs {
226 let name = d.function.name.clone();
227 let exec = exec.clone();
228 self.executor.add(
229 name.clone(),
230 Box::new(ClosureTool::boxed(move |s: String| {
231 let name = name.clone();
232 let exec = exec.clone();
233 Box::pin(async move { exec(name.clone(), s).await })
234 })),
235 );
236 }
237 self.tools.extend(defs);
238 self
239 }
240
241 /// Execute one iteration of the agent loop.
242 /// Returns `Ok(Some(content))` if loop should terminate, `Ok(None)` to continue
243 ///
244 /// This is usually used to customize the agent loop.
245 ///
246 /// # Example
247 /// ```
248 /// use tiny_loop::{Agent, llm::OpenAIProvider};
249 ///
250 /// # async fn example() -> tiny_loop::Result<()> {
251 /// let mut agent = Agent::new(OpenAIProvider::new())
252 /// .system("You are a helpful assistant");
253 ///
254 /// // Custom loop with early break
255 /// let mut iterations = 0;
256 /// loop {
257 /// println!("Before step {}", iterations);
258 ///
259 /// if let Some(content) = agent.step().await? {
260 /// println!("Completed: {}", content);
261 /// break;
262 /// }
263 ///
264 /// iterations += 1;
265 /// if iterations > 10 {
266 /// println!("Max iterations reached");
267 /// break;
268 /// }
269 /// }
270 /// # Ok(())
271 /// # }
272 /// ```
273 pub async fn step(&mut self) -> crate::Result<Option<String>> {
274 tracing::trace!("Calling LLM with {} messages", self.history.get_all().len());
275
276 let messages: Vec<_> = self
277 .history
278 .get_all()
279 .iter()
280 .map(|tm| tm.message.clone())
281 .collect();
282 let start = std::time::SystemTime::now();
283 let response = self.llm.call(&messages, &self.tools).await?;
284 let elapsed = start.elapsed().unwrap();
285
286 self.history.add(crate::types::TimedMessage {
287 message: response.message.clone().into(),
288 timestamp: start + elapsed,
289 elapsed,
290 });
291
292 // Execute tool calls if any
293 if let Some(calls) = &response.message.tool_calls {
294 tracing::debug!("Executing {} tool calls", calls.len());
295 let results = self.executor.execute(calls.clone()).await;
296 self.history.add_batch(
297 results
298 .into_iter()
299 .map(|r| crate::types::TimedMessage {
300 message: r.tool_message.into(),
301 timestamp: r.timestamp,
302 elapsed: r.elapsed,
303 })
304 .collect(),
305 );
306 }
307
308 // Break loop if finish reason is not tool_calls
309 if !matches!(
310 response.finish_reason,
311 crate::types::FinishReason::ToolCalls
312 ) {
313 tracing::debug!(
314 "Agent loop completed, finish_reason: {:?}",
315 response.finish_reason
316 );
317 return Ok(Some(response.message.content));
318 }
319
320 Ok(None)
321 }
322
323 /// Run the agent loop until completion.
324 /// Return the last AI's response
325 pub async fn run(&mut self) -> crate::Result<String> {
326 tracing::debug!("Starting agent loop");
327 loop {
328 if let Some(content) = self.step().await? {
329 return Ok(content);
330 }
331 }
332 }
333
334 /// Run the agent loop with a new user input appended.
335 /// Return the last AI's response
336 pub async fn chat(&mut self, prompt: impl Into<String>) -> crate::Result<String> {
337 let prompt = prompt.into();
338 tracing::debug!("Chat request, prompt length: {}", prompt.len());
339 self.history.add(crate::types::TimedMessage {
340 message: crate::types::UserMessage { content: prompt }.into(),
341 timestamp: std::time::SystemTime::now(),
342 elapsed: std::time::Duration::ZERO,
343 });
344 self.run().await
345 }
346}