Skip to main content

max / makenotwork

6.3 KB · 219 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::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::tool::{SurfaceProjection, ToolRegistry, ToolSpec};
20
21 /// A single turn in the conversation. Provider-neutral.
22 #[derive(Debug, Clone, Serialize, Deserialize)]
23 pub struct Message {
24 pub role: Role,
25 #[serde(default, skip_serializing_if = "Option::is_none")]
26 pub content: Option<String>,
27 #[serde(default, skip_serializing_if = "Vec::is_empty")]
28 pub tool_calls: Vec<ToolCall>,
29 /// When `role == Tool`, the name of the tool that produced this result.
30 #[serde(default, skip_serializing_if = "Option::is_none")]
31 pub tool_name: Option<String>,
32 }
33
34 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
35 #[serde(rename_all = "lowercase")]
36 pub enum Role {
37 System,
38 User,
39 Assistant,
40 Tool,
41 }
42
43 /// A tool invocation requested by the model.
44 #[derive(Debug, Clone, Serialize, Deserialize)]
45 pub struct ToolCall {
46 /// Provider-supplied ID, if any. Empty string if the provider doesn't emit one.
47 #[serde(default)]
48 pub id: String,
49 pub name: String,
50 #[serde(default)]
51 pub arguments: Value,
52 }
53
54 impl Message {
55 pub fn system(text: impl Into<String>) -> Self {
56 Self {
57 role: Role::System,
58 content: Some(text.into()),
59 tool_calls: Vec::new(),
60 tool_name: None,
61 }
62 }
63 pub fn user(text: impl Into<String>) -> Self {
64 Self {
65 role: Role::User,
66 content: Some(text.into()),
67 tool_calls: Vec::new(),
68 tool_name: None,
69 }
70 }
71 pub fn tool(name: impl Into<String>, content: impl Into<String>) -> Self {
72 Self {
73 role: Role::Tool,
74 content: Some(content.into()),
75 tool_calls: Vec::new(),
76 tool_name: Some(name.into()),
77 }
78 }
79 }
80
81 /// Provider abstraction: given a conversation and a tool surface, return the
82 /// next assistant message. Whether the message carries tool calls or a plain
83 /// text turn is up to the model.
84 #[async_trait]
85 pub trait InferenceProvider: Send + Sync {
86 async fn chat(&self, messages: &[Message], tools: &[ToolSpec]) -> Result<Message>;
87 }
88
89 /// Configuration for a [`Agent`] run.
90 pub struct AgentConfig {
91 pub max_steps: usize,
92 pub projection: SurfaceProjection,
93 /// Which write capabilities are granted for this session. `None` bypasses
94 /// the check entirely — appropriate only for fully-trusted callers.
95 pub grants: Option<HashSet<String>>,
96 }
97
98 impl Default for AgentConfig {
99 fn default() -> Self {
100 Self {
101 max_steps: 8,
102 projection: SurfaceProjection::Full,
103 grants: Some(HashSet::new()),
104 }
105 }
106 }
107
108 pub struct Agent<P: InferenceProvider> {
109 provider: P,
110 config: AgentConfig,
111 }
112
113 impl<P: InferenceProvider> Agent<P> {
114 pub fn new(provider: P) -> Self {
115 Self {
116 provider,
117 config: AgentConfig::default(),
118 }
119 }
120
121 #[must_use]
122 pub fn with_max_steps(mut self, n: usize) -> Self {
123 self.config.max_steps = n;
124 self
125 }
126
127 #[must_use]
128 pub fn with_projection(mut self, projection: SurfaceProjection) -> Self {
129 self.config.projection = projection;
130 self
131 }
132
133 #[must_use]
134 pub fn with_grants(mut self, grants: HashSet<String>) -> Self {
135 self.config.grants = Some(grants);
136 self
137 }
138
139 /// Bypass capability checks. Only appropriate for fully-trusted callers.
140 #[must_use]
141 pub fn without_grant_checks(mut self) -> Self {
142 self.config.grants = None;
143 self
144 }
145
146 /// Run the loop. Returns the final assistant message plus the transcript.
147 pub async fn run(
148 &self,
149 registry: &ToolRegistry,
150 system: Option<&str>,
151 user: &str,
152 ) -> Result<RunOutcome> {
153 let tools = registry.specs_projected(self.config.projection);
154
155 let mut messages: Vec<Message> = Vec::new();
156 if let Some(s) = system {
157 messages.push(Message::system(s));
158 }
159 messages.push(Message::user(user));
160
161 for step in 0..self.config.max_steps {
162 let assistant = self.provider.chat(&messages, &tools).await?;
163 let has_tool_calls = !assistant.tool_calls.is_empty();
164 messages.push(assistant.clone());
165
166 if !has_tool_calls {
167 return Ok(RunOutcome {
168 final_message: assistant.content.unwrap_or_default(),
169 transcript: messages,
170 steps: step + 1,
171 });
172 }
173
174 for call in assistant.tool_calls {
175 let result = registry
176 .call(&call.name, call.arguments, self.config.grants.as_ref())
177 .await;
178 let content = match result {
179 Ok(r) => tool_result_to_text(&r),
180 Err(e) => format!("[error] {e}"),
181 };
182 messages.push(Message::tool(&call.name, content));
183 }
184 }
185
186 Err(Error::Protocol(format!(
187 "agent exceeded max_steps ({})",
188 self.config.max_steps
189 )))
190 }
191 }
192
193 fn tool_result_to_text(result: &crate::tool::ToolCallResult) -> String {
194 use crate::tool::ContentPart;
195 let mut out = String::new();
196 for part in &result.content {
197 match part {
198 ContentPart::Text { text } => {
199 if !out.is_empty() {
200 out.push('\n');
201 }
202 out.push_str(text);
203 }
204 }
205 }
206 if result.is_error {
207 format!("[error] {out}")
208 } else {
209 out
210 }
211 }
212
213 /// Outcome of a completed agent run.
214 pub struct RunOutcome {
215 pub final_message: String,
216 pub transcript: Vec<Message>,
217 pub steps: usize,
218 }
219