1use super::{AgentRunTime, prompt::PromptComposer};
2use crate::{
3 checkpoint::CheckpointStorage,
4 error::AgentError,
5 messages::MessageChannel,
6 middleware::{Middleware, MiddlewareChain},
7 providers::LlmProvider,
8 sandbox::Sandbox,
9 skills::SkillManager,
10 tools::{ToolRegistry, register_skill_tools},
11 transcript::TranscriptStorage,
12};
13use std::{collections::HashMap, path::PathBuf, sync::Arc, time::Duration};
14use tokio::sync::Mutex;
15
16const DEFAULT_MAX_ITER: u32 = 10;
17const DEFAULT_MAX_CONSECUTIVE_FAIL_COUNT: u32 = 3;
18const DEFAULT_RETRY_BACKOFF: Duration = Duration::from_millis(200);
19const DEFAULT_TOOL_TIMEOUT: Duration = Duration::from_secs(120);
20
21#[derive(Clone)]
22pub struct AgentRuntimeConfig {
23 pub provider: Arc<dyn LlmProvider>,
24 pub model: String,
25 pub system_prompt: Option<String>,
26 pub temperature: Option<f32>,
27 pub max_tokens: Option<u32>,
28 pub max_iter: u32,
29 pub max_consecutive_fail_count: u32,
30 pub checkpoint_storage: Arc<dyn CheckpointStorage>,
31 pub transcript_storage: Arc<dyn TranscriptStorage>,
32 pub sandbox: Arc<dyn Sandbox>,
33 pub tools: ToolRegistry,
34 pub skill_dirs: Vec<PathBuf>,
35 pub middleware: MiddlewareChain,
36 pub messages: MessageChannel,
37 pub retry_backoff: Duration,
38 pub tool_timeout: Duration,
39}
40
41#[derive(Clone)]
42pub struct AgentRuntimeStorage {
43 pub checkpoint: Arc<dyn CheckpointStorage>,
44 pub transcript: Arc<dyn TranscriptStorage>,
45}
46
47impl AgentRuntimeStorage {
48 pub fn new(
49 checkpoint: Arc<dyn CheckpointStorage>,
50 transcript: Arc<dyn TranscriptStorage>,
51 ) -> Self {
52 Self {
53 checkpoint,
54 transcript,
55 }
56 }
57
58 pub fn shared<S>(storage: Arc<S>) -> Self
59 where
60 S: CheckpointStorage + TranscriptStorage + 'static,
61 {
62 Self {
63 checkpoint: storage.clone(),
64 transcript: storage,
65 }
66 }
67}
68
69#[derive(Clone)]
70pub struct AgentRuntimeOptions {
71 pub system_prompt: Option<String>,
72 pub temperature: Option<f32>,
73 pub max_tokens: Option<u32>,
74 pub max_iter: u32,
75 pub max_consecutive_fail_count: u32,
76 pub tools: ToolRegistry,
77 pub skill_dirs: Vec<PathBuf>,
78 pub middleware: MiddlewareChain,
79 pub messages: MessageChannel,
80 pub retry_backoff: Duration,
81 pub tool_timeout: Duration,
82}
83
84impl Default for AgentRuntimeOptions {
85 fn default() -> Self {
86 Self {
87 system_prompt: None,
88 temperature: None,
89 max_tokens: None,
90 max_iter: DEFAULT_MAX_ITER,
91 max_consecutive_fail_count: DEFAULT_MAX_CONSECUTIVE_FAIL_COUNT,
92 tools: ToolRegistry::new(),
93 skill_dirs: Vec::new(),
94 middleware: MiddlewareChain::new(),
95 messages: MessageChannel::disconnected(),
96 retry_backoff: DEFAULT_RETRY_BACKOFF,
97 tool_timeout: DEFAULT_TOOL_TIMEOUT,
98 }
99 }
100}
101
102#[derive(Clone)]
103pub struct AgentRuntimeConfigInput {
104 pub provider: Arc<dyn LlmProvider>,
105 pub model: String,
106 pub storage: AgentRuntimeStorage,
107 pub sandbox: Arc<dyn Sandbox>,
108 pub options: AgentRuntimeOptions,
109}
110
111impl AgentRuntimeConfig {
112 pub fn new(input: AgentRuntimeConfigInput) -> Self {
113 Self {
114 provider: input.provider,
115 model: input.model,
116 system_prompt: input.options.system_prompt,
117 temperature: input.options.temperature,
118 max_tokens: input.options.max_tokens,
119 max_iter: input.options.max_iter,
120 max_consecutive_fail_count: input.options.max_consecutive_fail_count,
121 checkpoint_storage: input.storage.checkpoint,
122 transcript_storage: input.storage.transcript,
123 sandbox: input.sandbox,
124 tools: input.options.tools,
125 skill_dirs: input.options.skill_dirs,
126 middleware: input.options.middleware,
127 messages: input.options.messages,
128 retry_backoff: input.options.retry_backoff,
129 tool_timeout: input.options.tool_timeout,
130 }
131 }
132}
133
134impl AgentRunTime {
135 pub fn from_config(config: AgentRuntimeConfig) -> Result<Self, AgentError> {
136 let mut tools = config.tools;
137 let skills = if config.skill_dirs.is_empty() {
138 None
139 } else {
140 Some(Arc::new(SkillManager::discover_dirs(&config.skill_dirs)?))
141 };
142 if let Some(skills) = &skills {
143 register_skill_tools(&mut tools, skills.clone());
144 }
145
146 Ok(Self {
147 provider: config.provider,
148 max_iter: config.max_iter,
149 max_consecutive_fail_count: config.max_consecutive_fail_count,
150 model: config.model,
151 system_prompt: config.system_prompt,
152 temperature: config.temperature,
153 max_tokens: config.max_tokens,
154 checkpoint_storage: config.checkpoint_storage,
155 transcript_storage: config.transcript_storage,
156 sandbox: config.sandbox,
157 tools,
158 skills,
159 prompt_composer: PromptComposer,
160 middleware: config.middleware,
161 messages: config.messages,
162 tasks: crate::tasks::TaskManager::new(),
163 retry_backoff: config.retry_backoff,
164 tool_timeout: config.tool_timeout,
165 session_locks: Mutex::new(HashMap::new()),
166 })
167 }
168}
169
170pub struct AgentRunTimeBuilder {
188 provider: Option<Arc<dyn LlmProvider>>,
189 model: Option<String>,
190 system_prompt: Option<String>,
191 temperature: Option<f32>,
192 max_tokens: Option<u32>,
193 max_iter: u32,
194 max_consecutive_fail_count: u32,
195 checkpoint_storage: Option<Arc<dyn CheckpointStorage>>,
196 transcript_storage: Option<Arc<dyn TranscriptStorage>>,
197 sandbox: Option<Arc<dyn Sandbox>>,
198 tools: ToolRegistry,
199 skill_dirs: Vec<PathBuf>,
200 middleware: MiddlewareChain,
201 messages: Option<MessageChannel>,
202 retry_backoff: Duration,
203 tool_timeout: Duration,
204}
205
206impl AgentRunTimeBuilder {
207 pub fn new() -> Self {
208 Self {
209 provider: None,
210 model: None,
211 system_prompt: None,
212 temperature: None,
213 max_tokens: None,
214 max_iter: DEFAULT_MAX_ITER,
215 max_consecutive_fail_count: DEFAULT_MAX_CONSECUTIVE_FAIL_COUNT,
216 checkpoint_storage: None,
217 transcript_storage: None,
218 sandbox: None,
219 tools: ToolRegistry::new(),
220 skill_dirs: Vec::new(),
221 middleware: MiddlewareChain::new(),
222 messages: None,
223 retry_backoff: DEFAULT_RETRY_BACKOFF,
224 tool_timeout: DEFAULT_TOOL_TIMEOUT,
225 }
226 }
227
228 pub fn provider(mut self, provider: Arc<dyn LlmProvider>) -> Self {
229 self.provider = Some(provider);
230 self
231 }
232
233 pub fn model(mut self, model: impl Into<String>) -> Self {
234 self.model = Some(model.into());
235 self
236 }
237
238 pub fn system_prompt(mut self, system_prompt: impl Into<String>) -> Self {
239 self.system_prompt = Some(system_prompt.into());
240 self
241 }
242
243 pub fn temperature(mut self, temperature: f32) -> Self {
244 self.temperature = Some(temperature);
245 self
246 }
247
248 pub fn max_tokens(mut self, max_tokens: u32) -> Self {
249 self.max_tokens = Some(max_tokens);
250 self
251 }
252
253 pub fn max_iter(mut self, max_iter: u32) -> Self {
254 self.max_iter = max_iter;
255 self
256 }
257
258 pub fn max_consecutive_fail_count(mut self, count: u32) -> Self {
259 self.max_consecutive_fail_count = count;
260 self
261 }
262
263 pub fn checkpoint_storage(mut self, storage: Arc<dyn CheckpointStorage>) -> Self {
264 self.checkpoint_storage = Some(storage);
265 self
266 }
267
268 pub fn transcript_storage(mut self, storage: Arc<dyn TranscriptStorage>) -> Self {
269 self.transcript_storage = Some(storage);
270 self
271 }
272
273 pub fn sandbox(mut self, sandbox: Arc<dyn Sandbox>) -> Self {
274 self.sandbox = Some(sandbox);
275 self
276 }
277
278 pub fn tools(mut self, tools: ToolRegistry) -> Self {
279 self.tools = tools;
280 self
281 }
282
283 pub fn skill_dir(mut self, path: impl Into<PathBuf>) -> Self {
284 self.skill_dirs.push(path.into());
285 self
286 }
287
288 pub fn skill_dirs<I, P>(mut self, paths: I) -> Self
289 where
290 I: IntoIterator<Item = P>,
291 P: Into<PathBuf>,
292 {
293 self.skill_dirs.extend(paths.into_iter().map(Into::into));
294 self
295 }
296
297 pub fn middleware(mut self, middleware: Arc<dyn Middleware>) -> Self {
299 self.middleware.push(middleware);
300 self
301 }
302
303 pub fn messages(mut self, messages: MessageChannel) -> Self {
305 self.messages = Some(messages);
306 self
307 }
308
309 pub fn retry_backoff(mut self, base: Duration) -> Self {
312 self.retry_backoff = base;
313 self
314 }
315
316 pub fn tool_timeout(mut self, timeout: Duration) -> Self {
318 self.tool_timeout = timeout;
319 self
320 }
321
322 pub fn build(self) -> Result<AgentRunTime, AgentError> {
323 AgentRunTime::from_config(AgentRuntimeConfig::new(AgentRuntimeConfigInput {
324 provider: self
325 .provider
326 .ok_or_else(|| AgentError::config("provider is required"))?,
327 model: self
328 .model
329 .ok_or_else(|| AgentError::config("model is required"))?,
330 storage: AgentRuntimeStorage::new(
331 self.checkpoint_storage
332 .ok_or_else(|| AgentError::config("checkpoint_storage is required"))?,
333 self.transcript_storage
334 .ok_or_else(|| AgentError::config("transcript_storage is required"))?,
335 ),
336 sandbox: self
337 .sandbox
338 .ok_or_else(|| AgentError::config("sandbox is required"))?,
339 options: AgentRuntimeOptions {
340 system_prompt: self.system_prompt,
341 temperature: self.temperature,
342 max_tokens: self.max_tokens,
343 max_iter: self.max_iter,
344 max_consecutive_fail_count: self.max_consecutive_fail_count,
345 tools: self.tools,
346 skill_dirs: self.skill_dirs,
347 middleware: self.middleware,
348 messages: self.messages.unwrap_or_default(),
349 retry_backoff: self.retry_backoff,
350 tool_timeout: self.tool_timeout,
351 },
352 }))
353 }
354}
355
356impl Default for AgentRunTimeBuilder {
357 fn default() -> Self {
358 Self::new()
359 }
360}