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