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