tiny_agent/core/
builder.rs

1use super::{AgentRunTime, prompt::PromptComposer};
2use crate::{
3    checkpoint::CheckpointStorage,
4    error::AgentError,
5    middleware::{Middleware, MiddlewareChain},
6    providers::LlmProvider,
7    sandbox::Sandbox,
8    skills::SkillManager,
9    tools::{ToolRegistry, register_skill_tools},
10    trajectory::{NoopSink, TrajectorySink},
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 trajectory: Arc<dyn TrajectorySink>,
37    pub retry_backoff: Duration,
38    pub tool_timeout: Duration,
39}
40
41impl AgentRuntimeConfig {
42    pub fn new(
43        provider: Arc<dyn LlmProvider>,
44        model: impl Into<String>,
45        checkpoint_storage: Arc<dyn CheckpointStorage>,
46        transcript_storage: Arc<dyn TranscriptStorage>,
47        sandbox: Arc<dyn Sandbox>,
48    ) -> Self {
49        Self {
50            provider,
51            model: model.into(),
52            system_prompt: None,
53            temperature: None,
54            max_tokens: None,
55            max_iter: DEFAULT_MAX_ITER,
56            max_consecutive_fail_count: DEFAULT_MAX_CONSECUTIVE_FAIL_COUNT,
57            checkpoint_storage,
58            transcript_storage,
59            sandbox,
60            tools: ToolRegistry::new(),
61            skill_dirs: Vec::new(),
62            middleware: MiddlewareChain::new(),
63            trajectory: Arc::new(NoopSink),
64            retry_backoff: DEFAULT_RETRY_BACKOFF,
65            tool_timeout: DEFAULT_TOOL_TIMEOUT,
66        }
67    }
68
69    pub fn with_tools(mut self, tools: ToolRegistry) -> Self {
70        self.tools = tools;
71        self
72    }
73
74    pub fn with_system_prompt(mut self, system_prompt: impl Into<String>) -> Self {
75        self.system_prompt = Some(system_prompt.into());
76        self
77    }
78
79    pub fn with_middleware(mut self, middleware: Arc<dyn Middleware>) -> Self {
80        self.middleware.push(middleware);
81        self
82    }
83
84    pub fn with_trajectory(mut self, trajectory: Arc<dyn TrajectorySink>) -> Self {
85        self.trajectory = trajectory;
86        self
87    }
88}
89
90impl AgentRunTime {
91    pub fn from_config(config: AgentRuntimeConfig) -> Result<Self, AgentError> {
92        let mut tools = config.tools;
93        let skills = if config.skill_dirs.is_empty() {
94            None
95        } else {
96            Some(Arc::new(SkillManager::discover_dirs(&config.skill_dirs)?))
97        };
98        if let Some(skills) = &skills {
99            register_skill_tools(&mut tools, skills.clone());
100        }
101
102        Ok(Self {
103            provider: config.provider,
104            max_iter: config.max_iter,
105            max_consecutive_fail_count: config.max_consecutive_fail_count,
106            model: config.model,
107            system_prompt: config.system_prompt,
108            temperature: config.temperature,
109            max_tokens: config.max_tokens,
110            checkpoint_storage: config.checkpoint_storage,
111            transcript_storage: config.transcript_storage,
112            sandbox: config.sandbox,
113            tools,
114            skills,
115            prompt_composer: PromptComposer,
116            middleware: config.middleware,
117            trajectory: config.trajectory,
118            retry_backoff: config.retry_backoff,
119            tool_timeout: config.tool_timeout,
120            session_locks: Mutex::new(HashMap::new()),
121        })
122    }
123}
124
125/// `AgentRunTime` 的构造器。
126///
127/// `provider` / `model` / 三个存储 / `sandbox` 是必填项;
128/// 中间件与轨迹 sink 可选(默认空链 + `NoopSink`)。
129///
130/// ```ignore
131/// let runtime = AgentRunTimeBuilder::new()
132///     .provider(provider)
133///     .model("gpt-4o")
134///     .checkpoint_storage(storage.clone())
135///     .transcript_storage(storage)
136///     .sandbox(sandbox)
137///     .tools(tools)
138///     .middleware(Arc::new(MyGuardrail))
139///     .trajectory(sink)
140///     .build()?;
141/// ```
142pub struct AgentRunTimeBuilder {
143    provider: Option<Arc<dyn LlmProvider>>,
144    model: Option<String>,
145    system_prompt: Option<String>,
146    temperature: Option<f32>,
147    max_tokens: Option<u32>,
148    max_iter: u32,
149    max_consecutive_fail_count: u32,
150    checkpoint_storage: Option<Arc<dyn CheckpointStorage>>,
151    transcript_storage: Option<Arc<dyn TranscriptStorage>>,
152    sandbox: Option<Arc<dyn Sandbox>>,
153    tools: ToolRegistry,
154    skill_dirs: Vec<PathBuf>,
155    middleware: MiddlewareChain,
156    trajectory: Option<Arc<dyn TrajectorySink>>,
157    retry_backoff: Duration,
158    tool_timeout: Duration,
159}
160
161impl AgentRunTimeBuilder {
162    pub fn new() -> Self {
163        Self {
164            provider: None,
165            model: None,
166            system_prompt: None,
167            temperature: None,
168            max_tokens: None,
169            max_iter: DEFAULT_MAX_ITER,
170            max_consecutive_fail_count: DEFAULT_MAX_CONSECUTIVE_FAIL_COUNT,
171            checkpoint_storage: None,
172            transcript_storage: None,
173            sandbox: None,
174            tools: ToolRegistry::new(),
175            skill_dirs: Vec::new(),
176            middleware: MiddlewareChain::new(),
177            trajectory: None,
178            retry_backoff: DEFAULT_RETRY_BACKOFF,
179            tool_timeout: DEFAULT_TOOL_TIMEOUT,
180        }
181    }
182
183    pub fn provider(mut self, provider: Arc<dyn LlmProvider>) -> Self {
184        self.provider = Some(provider);
185        self
186    }
187
188    pub fn model(mut self, model: impl Into<String>) -> Self {
189        self.model = Some(model.into());
190        self
191    }
192
193    pub fn system_prompt(mut self, system_prompt: impl Into<String>) -> Self {
194        self.system_prompt = Some(system_prompt.into());
195        self
196    }
197
198    pub fn with_system_prompt(self, system_prompt: impl Into<String>) -> Self {
199        self.system_prompt(system_prompt)
200    }
201
202    pub fn temperature(mut self, temperature: f32) -> Self {
203        self.temperature = Some(temperature);
204        self
205    }
206
207    pub fn max_tokens(mut self, max_tokens: u32) -> Self {
208        self.max_tokens = Some(max_tokens);
209        self
210    }
211
212    pub fn max_iter(mut self, max_iter: u32) -> Self {
213        self.max_iter = max_iter;
214        self
215    }
216
217    pub fn max_consecutive_fail_count(mut self, count: u32) -> Self {
218        self.max_consecutive_fail_count = count;
219        self
220    }
221
222    pub fn checkpoint_storage(mut self, storage: Arc<dyn CheckpointStorage>) -> Self {
223        self.checkpoint_storage = Some(storage);
224        self
225    }
226
227    pub fn transcript_storage(mut self, storage: Arc<dyn TranscriptStorage>) -> Self {
228        self.transcript_storage = Some(storage);
229        self
230    }
231
232    pub fn sandbox(mut self, sandbox: Arc<dyn Sandbox>) -> Self {
233        self.sandbox = Some(sandbox);
234        self
235    }
236
237    pub fn tools(mut self, tools: ToolRegistry) -> Self {
238        self.tools = tools;
239        self
240    }
241
242    pub fn skill_dir(mut self, path: impl Into<PathBuf>) -> Self {
243        self.skill_dirs.push(path.into());
244        self
245    }
246
247    pub fn skill_dirs<I, P>(mut self, paths: I) -> Self
248    where
249        I: IntoIterator<Item = P>,
250        P: Into<PathBuf>,
251    {
252        self.skill_dirs.extend(paths.into_iter().map(Into::into));
253        self
254    }
255
256    /// 追加一个中间件(可多次调用,按调用顺序构成 onion 链)。
257    pub fn middleware(mut self, middleware: Arc<dyn Middleware>) -> Self {
258        self.middleware.push(middleware);
259        self
260    }
261
262    /// 设置轨迹 sink(不设则为 `NoopSink`)。
263    pub fn trajectory(mut self, trajectory: Arc<dyn TrajectorySink>) -> Self {
264        self.trajectory = Some(trajectory);
265        self
266    }
267
268    /// 可重试失败的退避基数(默认 200ms,实际等待为指数退避 + jitter)。
269    /// 传 `Duration::ZERO` 可关闭退避(测试常用)。
270    pub fn retry_backoff(mut self, base: Duration) -> Self {
271        self.retry_backoff = base;
272        self
273    }
274
275    /// 单个工具调用的超时时间(默认 120s)。传 `Duration::ZERO` 可关闭工具超时。
276    pub fn tool_timeout(mut self, timeout: Duration) -> Self {
277        self.tool_timeout = timeout;
278        self
279    }
280
281    pub fn build(self) -> Result<AgentRunTime, AgentError> {
282        AgentRunTime::from_config(AgentRuntimeConfig {
283            provider: self
284                .provider
285                .ok_or_else(|| AgentError::config("provider is required"))?,
286            model: self
287                .model
288                .ok_or_else(|| AgentError::config("model is required"))?,
289            system_prompt: self.system_prompt,
290            temperature: self.temperature,
291            max_tokens: self.max_tokens,
292            max_iter: self.max_iter,
293            max_consecutive_fail_count: self.max_consecutive_fail_count,
294            checkpoint_storage: self
295                .checkpoint_storage
296                .ok_or_else(|| AgentError::config("checkpoint_storage is required"))?,
297            transcript_storage: self
298                .transcript_storage
299                .ok_or_else(|| AgentError::config("transcript_storage is required"))?,
300            sandbox: self
301                .sandbox
302                .ok_or_else(|| AgentError::config("sandbox is required"))?,
303            tools: self.tools,
304            skill_dirs: self.skill_dirs,
305            middleware: self.middleware,
306            trajectory: self.trajectory.unwrap_or_else(|| Arc::new(NoopSink)),
307            retry_backoff: self.retry_backoff,
308            tool_timeout: self.tool_timeout,
309        })
310    }
311}
312
313impl Default for AgentRunTimeBuilder {
314    fn default() -> Self {
315        Self::new()
316    }
317}
tiny_agent/core/builder.rs

tiny_agent/core/
builder.rs
