Skip to main content

max / makenotwork

12.9 KB · 421 lines History Blame Raw
1 //! Provider-agnostic agent loop.
2 //!
3 //! The [`Agent`] runs the tool-use loop: send `messages + tools` to an
4 //! [`InferenceProvider`], dispatch any returned tool calls against a
5 //! [`ToolRegistry`], feed the results back, repeat until the provider returns
6 //! a plain assistant message or the step budget is exhausted.
7 //!
8 //! Providers (Ollama, and later OpenAI / Anthropic / OpenAI-compatible) only
9 //! need to implement the small [`InferenceProvider`] trait; the loop is
10 //! reused across all of them.
11
12 use std::collections::{HashMap, HashSet};
13
14 use async_trait::async_trait;
15 use serde::{Deserialize, Serialize};
16 use serde_json::Value;
17
18 use crate::error::{Error, Result};
19 use crate::schema::render_slots;
20 use crate::tool::{SurfaceProjection, ToolRegistry, ToolSpec};
21
22 /// A single turn in the conversation. Provider-neutral.
23 #[derive(Debug, Clone, Serialize, Deserialize)]
24 pub struct Message {
25 pub role: Role,
26 #[serde(default, skip_serializing_if = "Option::is_none")]
27 pub content: Option<String>,
28 #[serde(default, skip_serializing_if = "Vec::is_empty")]
29 pub tool_calls: Vec<ToolCall>,
30 /// When `role == Tool`, the name of the tool that produced this result.
31 #[serde(default, skip_serializing_if = "Option::is_none")]
32 pub tool_name: Option<String>,
33 }
34
35 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
36 #[serde(rename_all = "lowercase")]
37 pub enum Role {
38 System,
39 User,
40 Assistant,
41 Tool,
42 }
43
44 /// A tool invocation requested by the model.
45 #[derive(Debug, Clone, Serialize, Deserialize)]
46 pub struct ToolCall {
47 /// Provider-supplied ID, if any. Empty string if the provider doesn't emit one.
48 #[serde(default)]
49 pub id: String,
50 pub name: String,
51 #[serde(default)]
52 pub arguments: Value,
53 }
54
55 impl Message {
56 pub fn system(text: impl Into<String>) -> Self {
57 Self {
58 role: Role::System,
59 content: Some(text.into()),
60 tool_calls: Vec::new(),
61 tool_name: None,
62 }
63 }
64 pub fn user(text: impl Into<String>) -> Self {
65 Self {
66 role: Role::User,
67 content: Some(text.into()),
68 tool_calls: Vec::new(),
69 tool_name: None,
70 }
71 }
72 pub fn tool(name: impl Into<String>, content: impl Into<String>) -> Self {
73 Self {
74 role: Role::Tool,
75 content: Some(content.into()),
76 tool_calls: Vec::new(),
77 tool_name: Some(name.into()),
78 }
79 }
80 }
81
82 /// Provider abstraction: given a conversation and a tool surface, return the
83 /// next assistant message. Whether the message carries tool calls or a plain
84 /// text turn is up to the model.
85 #[async_trait]
86 pub trait InferenceProvider: Send + Sync {
87 async fn chat(&self, messages: &[Message], tools: &[ToolSpec]) -> Result<Message>;
88 }
89
90 /// Configuration for a [`Agent`] run.
91 pub struct AgentConfig {
92 pub max_steps: usize,
93 pub projection: SurfaceProjection,
94 /// Which write capabilities are granted for this session. `None` bypasses
95 /// the check entirely — appropriate only for fully-trusted callers.
96 pub grants: Option<HashSet<String>>,
97 /// How many times one tool may fail schema validation and still be invited
98 /// to try again. Past this, the diagnostic carries a stable "do not retry"
99 /// suffix so a model that cannot get the arguments right stops burning the
100 /// step budget on the same call.
101 pub max_repairs: usize,
102 }
103
104 impl Default for AgentConfig {
105 fn default() -> Self {
106 Self {
107 max_steps: 8,
108 projection: SurfaceProjection::Full,
109 grants: Some(HashSet::new()),
110 max_repairs: 2,
111 }
112 }
113 }
114
115 pub struct Agent<P: InferenceProvider> {
116 provider: P,
117 config: AgentConfig,
118 }
119
120 impl<P: InferenceProvider> Agent<P> {
121 pub fn new(provider: P) -> Self {
122 Self {
123 provider,
124 config: AgentConfig::default(),
125 }
126 }
127
128 #[must_use]
129 pub fn with_max_steps(mut self, n: usize) -> Self {
130 self.config.max_steps = n;
131 self
132 }
133
134 #[must_use]
135 pub fn with_projection(mut self, projection: SurfaceProjection) -> Self {
136 self.config.projection = projection;
137 self
138 }
139
140 #[must_use]
141 pub fn with_grants(mut self, grants: HashSet<String>) -> Self {
142 self.config.grants = Some(grants);
143 self
144 }
145
146 /// Bypass capability checks. Only appropriate for fully-trusted callers.
147 #[must_use]
148 pub fn without_grant_checks(mut self) -> Self {
149 self.config.grants = None;
150 self
151 }
152
153 #[must_use]
154 pub fn with_max_repairs(mut self, n: usize) -> Self {
155 self.config.max_repairs = n;
156 self
157 }
158
159 /// Run the loop. Returns the final assistant message plus the transcript.
160 pub async fn run(
161 &self,
162 registry: &ToolRegistry,
163 system: Option<&str>,
164 user: &str,
165 ) -> Result<RunOutcome> {
166 let tools = registry.specs_projected(self.config.projection);
167
168 let mut messages: Vec<Message> = Vec::new();
169 if let Some(s) = system {
170 messages.push(Message::system(s));
171 }
172 messages.push(Message::user(user));
173
174 // Per-tool count of schema-validation failures, so a model looping on
175 // the same malformed call can be told to stop rather than being
176 // invited to repair forever.
177 let mut repairs: HashMap<String, usize> = HashMap::new();
178
179 for step in 0..self.config.max_steps {
180 let assistant = self.provider.chat(&messages, &tools).await?;
181 let has_tool_calls = !assistant.tool_calls.is_empty();
182 messages.push(assistant.clone());
183
184 if !has_tool_calls {
185 return Ok(RunOutcome {
186 final_message: assistant.content.unwrap_or_default(),
187 transcript: messages,
188 steps: step + 1,
189 });
190 }
191
192 for call in assistant.tool_calls {
193 let result = registry
194 .call(&call.name, call.arguments, self.config.grants.as_ref())
195 .await;
196 let content = match result {
197 Ok(r) => tool_result_to_text(&r),
198 Err(Error::InvalidArguments { tool, slots }) => {
199 let seen = repairs.entry(tool.clone()).or_insert(0);
200 *seen += 1;
201 self.repair_prompt(&tool, &render_slots(&slots), *seen)
202 }
203 Err(e) => format!("[error] {e}"),
204 };
205 messages.push(Message::tool(&call.name, content));
206 }
207 }
208
209 Err(Error::Protocol(format!(
210 "agent exceeded max_steps ({})",
211 self.config.max_steps
212 )))
213 }
214
215 /// The tool message a model sees after failing schema validation.
216 ///
217 /// Names the offending fields and the tool to call again, because a model
218 /// handed only "invalid arguments" re-sends the same call. Past
219 /// `max_repairs` it flips to a stable refusal instead, matching the "do
220 /// not retry" suffix [`Error::CapabilityDenied`] uses.
221 fn repair_prompt(&self, tool: &str, slots: &str, attempt: usize) -> String {
222 if attempt > self.config.max_repairs {
223 format!(
224 "[error] invalid arguments for `{tool}`: {slots}. \
225 This call has now failed validation {attempt} times; \
226 do not retry `{tool}`."
227 )
228 } else {
229 format!(
230 "[error] invalid arguments for `{tool}`: {slots}. \
231 Correct the named fields and call `{tool}` again."
232 )
233 }
234 }
235 }
236
237 fn tool_result_to_text(result: &crate::tool::ToolCallResult) -> String {
238 use crate::tool::ContentPart;
239 let mut out = String::new();
240 for part in &result.content {
241 match part {
242 ContentPart::Text { text } => {
243 if !out.is_empty() {
244 out.push('\n');
245 }
246 out.push_str(text);
247 }
248 }
249 }
250 if result.is_error {
251 format!("[error] {out}")
252 } else {
253 out
254 }
255 }
256
257 /// Outcome of a completed agent run.
258 pub struct RunOutcome {
259 pub final_message: String,
260 pub transcript: Vec<Message>,
261 pub steps: usize,
262 }
263
264 #[cfg(test)]
265 mod tests {
266 use super::*;
267 use crate::tool::{Tool, ToolCallResult, ToolKind};
268 use serde_json::json;
269 use std::sync::Mutex;
270
271 /// Replays a fixed script of assistant turns, recording what it was told.
272 struct Scripted {
273 turns: Mutex<std::vec::IntoIter<Message>>,
274 seen: Mutex<Vec<Message>>,
275 }
276
277 impl Scripted {
278 fn new(turns: Vec<Message>) -> Self {
279 Self {
280 turns: Mutex::new(turns.into_iter()),
281 seen: Mutex::new(Vec::new()),
282 }
283 }
284
285 /// Every tool-role message the loop fed back.
286 fn tool_messages(&self) -> Vec<String> {
287 self.seen
288 .lock()
289 .unwrap()
290 .iter()
291 .filter(|m| m.role == Role::Tool)
292 .filter_map(|m| m.content.clone())
293 .collect()
294 }
295 }
296
297 #[async_trait]
298 impl InferenceProvider for Scripted {
299 async fn chat(&self, messages: &[Message], _tools: &[ToolSpec]) -> Result<Message> {
300 *self.seen.lock().unwrap() = messages.to_vec();
301 self.turns
302 .lock()
303 .unwrap()
304 .next()
305 .ok_or_else(|| Error::Protocol("script exhausted".into()))
306 }
307 }
308
309 fn call(name: &str, args: Value) -> Message {
310 Message {
311 role: Role::Assistant,
312 content: None,
313 tool_calls: vec![ToolCall {
314 id: String::new(),
315 name: name.to_string(),
316 arguments: args,
317 }],
318 tool_name: None,
319 }
320 }
321
322 fn done(text: &str) -> Message {
323 Message {
324 role: Role::Assistant,
325 content: Some(text.to_string()),
326 tool_calls: Vec::new(),
327 tool_name: None,
328 }
329 }
330
331 struct Greet;
332 #[async_trait]
333 impl Tool for Greet {
334 fn name(&self) -> &'static str {
335 "greet"
336 }
337 fn description(&self) -> &'static str {
338 "greets a name"
339 }
340 fn kind(&self) -> ToolKind {
341 ToolKind::Read
342 }
343 fn input_schema(&self) -> Value {
344 json!({
345 "type": "object",
346 "properties": { "name": { "type": "string" } },
347 "required": ["name"]
348 })
349 }
350 async fn call(&self, args: Value) -> Result<ToolCallResult> {
351 Ok(ToolCallResult::text(format!(
352 "hello {}",
353 args["name"].as_str().unwrap_or("?")
354 )))
355 }
356 }
357
358 fn registry() -> ToolRegistry {
359 let mut r = ToolRegistry::new();
360 r.register(Greet);
361 r
362 }
363
364 #[tokio::test]
365 async fn a_bad_call_is_repaired_on_the_next_turn() {
366 let provider = Scripted::new(vec![
367 call("greet", json!({})), // forgets `name`
368 call("greet", json!({ "name": "max" })), // corrects it
369 done("greeted"),
370 ]);
371 let agent = Agent::new(provider);
372 let out = agent.run(&registry(), None, "say hi").await.unwrap();
373
374 assert_eq!(out.final_message, "greeted");
375 assert_eq!(out.steps, 3);
376
377 let fed = agent.provider.tool_messages();
378 assert!(
379 fed[0].contains("missing required field `name` (string)"),
380 "diagnostic must name the field: {}",
381 fed[0]
382 );
383 assert!(
384 fed[0].contains("call `greet` again"),
385 "diagnostic must invite the retry: {}",
386 fed[0]
387 );
388 assert_eq!(fed[1], "hello max");
389 }
390
391 #[tokio::test]
392 async fn a_model_looping_on_the_same_bad_call_is_told_to_stop() {
393 let provider = Scripted::new(vec![
394 call("greet", json!({})),
395 call("greet", json!({})),
396 call("greet", json!({})),
397 done("gave up"),
398 ]);
399 let agent = Agent::new(provider).with_max_repairs(2);
400 let out = agent.run(&registry(), None, "say hi").await.unwrap();
401 assert_eq!(out.final_message, "gave up");
402
403 let fed = agent.provider.tool_messages();
404 assert!(fed[0].contains("call `greet` again"));
405 assert!(fed[1].contains("call `greet` again"));
406 assert!(
407 fed[2].contains("do not retry `greet`"),
408 "third failure must carry the stable refusal: {}",
409 fed[2]
410 );
411 }
412
413 #[tokio::test]
414 async fn a_valid_call_never_produces_a_diagnostic() {
415 let provider = Scripted::new(vec![call("greet", json!({ "name": "max" })), done("ok")]);
416 let agent = Agent::new(provider);
417 agent.run(&registry(), None, "say hi").await.unwrap();
418 assert_eq!(agent.provider.tool_messages(), ["hello max"]);
419 }
420 }
421