Skip to main content

tiny_agent/core/
builder.rs

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
170/// `AgentRunTime` 的构造器。
171///
172/// `provider` / `model` / 两类存储 / `sandbox` 是必填项;
173/// 中间件与实时消息通道可选。可观测性由标准 `tracing` span/event 产出,
174/// 是否导出到 OpenTelemetry / 日志 / 文件由宿主程序配置 subscriber 决定。
175///
176/// ```ignore
177/// let runtime = AgentRunTimeBuilder::new()
178///     .provider(provider)
179///     .model("gpt-4o")
180///     .checkpoint_storage(storage.clone())
181///     .transcript_storage(storage)
182///     .sandbox(sandbox)
183///     .tools(tools)
184///     .middleware(Arc::new(MyGuardrail))
185///     .build()?;
186/// ```
187pub 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    /// 追加一个中间件(可多次调用,按调用顺序构成 onion 链)。
298    pub fn middleware(mut self, middleware: Arc<dyn Middleware>) -> Self {
299        self.middleware.push(middleware);
300        self
301    }
302
303    /// 设置对前端的实时消息通道(不设则为 disconnected,消息静默丢弃)。
304    pub fn messages(mut self, messages: MessageChannel) -> Self {
305        self.messages = Some(messages);
306        self
307    }
308
309    /// 可重试失败的退避基数(默认 200ms,实际等待为指数退避 + jitter)。
310    /// 传 `Duration::ZERO` 可关闭退避(测试常用)。
311    pub fn retry_backoff(mut self, base: Duration) -> Self {
312        self.retry_backoff = base;
313        self
314    }
315
316    /// 单个工具调用的超时时间(默认 120s)。传 `Duration::ZERO` 可关闭工具超时。
317    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}