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