1use std::collections::HashSet;
2use std::sync::Arc;
3
4use agent_base::{AgentResult, AgentRuntime, LlmClient, Tool};
5
6#[cfg(feature = "skill")]
7use crate::skill::{LazySkillPrompter, Skill, SkillDetailTool, SkillPrompter};
8
9pub struct AgentBuilder {
10 inner: agent_base::AgentBuilder,
11 system_prompt: Option<String>,
12 tool_names: HashSet<String>,
13 #[cfg(feature = "skill")]
14 skills: Vec<Arc<dyn Skill>>,
15 #[cfg(feature = "skill")]
16 skill_prompter: Option<Arc<dyn SkillPrompter>>,
17 #[cfg(feature = "skill")]
18 skill_detail_tool_name: String,
19 #[cfg(feature = "skill")]
20 disable_skill_prompt_injection: bool,
21}
22
23impl AgentBuilder {
24 pub fn new(client: Arc<dyn LlmClient>) -> Self {
25 Self {
26 inner: agent_base::AgentBuilder::new(client),
27 system_prompt: None,
28 tool_names: HashSet::new(),
29 #[cfg(feature = "skill")]
30 skills: Vec::new(),
31 #[cfg(feature = "skill")]
32 skill_prompter: None,
33 #[cfg(feature = "skill")]
34 skill_detail_tool_name: "get_skill_detail".to_string(),
35 #[cfg(feature = "skill")]
36 disable_skill_prompt_injection: false,
37 }
38 }
39
40 pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
41 let prompt = prompt.into();
42 self.inner = self.inner.system_prompt(prompt.clone());
43 self.system_prompt = Some(prompt);
44 self
45 }
46
47 pub fn enable_thought(self, enable: bool) -> Self {
48 Self {
49 inner: self.inner.enable_thought(enable),
50 ..self
51 }
52 }
53
54 pub fn reasoning(self, config: agent_base::ReasoningConfig) -> Self {
55 Self {
56 inner: self.inner.reasoning(config),
57 ..self
58 }
59 }
60
61 pub fn enable_thinking(self, enable: bool) -> Self {
62 Self {
63 inner: self.inner.enable_thinking(enable),
64 ..self
65 }
66 }
67
68 pub fn thinking_budget(self, budget: u64) -> Self {
69 Self {
70 inner: self.inner.thinking_budget(budget),
71 ..self
72 }
73 }
74
75 pub fn tool_timeout(self, timeout_ms: u64) -> Self {
76 Self {
77 inner: self.inner.tool_timeout(timeout_ms),
78 ..self
79 }
80 }
81
82 pub fn max_tool_output_chars(self, max_chars: usize) -> Self {
83 Self {
84 inner: self.inner.max_tool_output_chars(max_chars),
85 ..self
86 }
87 }
88
89 pub fn register_tool(mut self, tool: impl Tool + 'static) -> Self {
90 self.tool_names.insert(tool.name().to_string());
91 self.inner = self.inner.register_tool(tool);
92 self
93 }
94
95 pub fn register_tool_arc(mut self, tool: Arc<dyn Tool>) -> Self {
96 self.tool_names.insert(tool.name().to_string());
97 self.inner = self.inner.register_tool_arc(tool);
98 self
99 }
100
101 pub fn approval_handler(self, handler: Arc<dyn agent_base::ApprovalHandler>) -> Self {
102 Self {
103 inner: self.inner.approval_handler(handler),
104 ..self
105 }
106 }
107
108 pub fn tool_policy(self, policy: Arc<dyn agent_base::ToolPolicy>) -> Self {
109 Self {
110 inner: self.inner.tool_policy(policy),
111 ..self
112 }
113 }
114
115 pub fn middleware(self, mw: impl agent_base::Middleware + 'static) -> Self {
116 Self {
117 inner: self.inner.middleware(mw),
118 ..self
119 }
120 }
121
122 pub fn context_window(self, max_tokens: usize) -> Self {
123 Self {
124 inner: self.inner.context_window(max_tokens),
125 ..self
126 }
127 }
128
129 pub fn context_window_manager(self, manager: agent_base::ContextWindowManager) -> Self {
130 Self {
131 inner: self.inner.context_window_manager(manager),
132 ..self
133 }
134 }
135
136 pub fn response_format(self, format: agent_base::ResponseFormat) -> Self {
137 Self {
138 inner: self.inner.response_format(format),
139 ..self
140 }
141 }
142
143 pub fn llm_retry(self, retry: agent_base::RetryConfig) -> Self {
144 Self {
145 inner: self.inner.llm_retry(retry),
146 ..self
147 }
148 }
149
150 pub fn session_store(self, store: Arc<dyn agent_base::SessionStore>) -> Self {
151 Self {
152 inner: self.inner.session_store(store),
153 ..self
154 }
155 }
156
157 pub fn error_recovery(self, recovery: Arc<dyn agent_base::ToolErrorRecovery>) -> Self {
158 Self {
159 inner: self.inner.error_recovery(recovery),
160 ..self
161 }
162 }
163
164 pub fn tool_error_retry_prompt(self, prompt: impl Into<String>) -> Self {
165 Self {
166 inner: self.inner.tool_error_retry_prompt(prompt),
167 ..self
168 }
169 }
170
171 pub fn language(self, language: agent_base::Language) -> Self {
172 Self {
173 inner: self.inner.language(language),
174 ..self
175 }
176 }
177
178 pub fn event_bus_capacity(self, capacity: usize) -> Self {
179 Self {
180 inner: self.inner.event_bus_capacity(capacity),
181 ..self
182 }
183 }
184
185 pub fn session_id_generator(
186 self,
187 generator: Arc<dyn agent_base::types::SessionIdGenerator>,
188 ) -> Self {
189 Self {
190 inner: self.inner.session_id_generator(generator),
191 ..self
192 }
193 }
194
195 #[cfg(feature = "skill")]
196 pub fn register_skill(mut self, skill: impl Skill + 'static) -> Self {
197 self.skills.push(Arc::new(skill));
198 self
199 }
200
201 #[cfg(feature = "skill")]
202 pub fn register_skills(mut self, skills: Vec<Arc<dyn Skill>>) -> Self {
203 self.skills.extend(skills);
204 self
205 }
206
207 #[cfg(feature = "skill")]
208 pub fn skill_prompter(mut self, prompter: Arc<dyn SkillPrompter>) -> Self {
209 self.skill_prompter = Some(prompter);
210 self
211 }
212
213 #[cfg(feature = "skill")]
214 pub fn disable_skill_prompt_injection(mut self) -> Self {
215 self.disable_skill_prompt_injection = true;
216 self
217 }
218
219 #[cfg(feature = "skill")]
220 pub fn skill_detail_tool_name(mut self, name: impl Into<String>) -> Self {
221 self.skill_detail_tool_name = name.into();
222 self
223 }
224
225 pub fn build(self) -> AgentResult<AgentRuntime> {
226 #[cfg(feature = "skill")]
227 {
228 self.build_with_skills()
229 }
230 #[cfg(not(feature = "skill"))]
231 {
232 self.inner.build()
233 }
234 }
235
236 #[cfg(feature = "skill")]
237 fn build_with_skills(mut self) -> AgentResult<AgentRuntime> {
238 let mut ab = self.inner;
239
240 if self.skills.is_empty() {
241 return ab.build();
242 }
243
244 let prompter: Arc<dyn SkillPrompter> = self
245 .skill_prompter
246 .take()
247 .unwrap_or_else(|| Arc::new(LazySkillPrompter::new()));
248
249 let mut skill_refs: Vec<Arc<dyn Skill>> = Vec::new();
250
251 for skill in self.skills {
252 for tool in skill.tools() {
253 let tool_name = tool.name().to_string();
254 if self.tool_names.contains(&tool_name) {
255 return Err(agent_base::AgentError::internal(format!(
256 "Tool name conflict: `{}` (Skill `{}`)",
257 tool_name,
258 skill.name()
259 )));
260 }
261 self.tool_names.insert(tool_name);
262 ab = ab.register_tool_arc(tool);
263 }
264 skill_refs.push(skill);
265 }
266
267 if !self.disable_skill_prompt_injection {
268 let skill_prompt = prompter.build_prompt(&skill_refs, &self.skill_detail_tool_name);
269 if !skill_prompt.is_empty() {
270 let new_prompt = match self.system_prompt.take() {
271 Some(existing) => format!("{}\n\n---\n\n{}", existing, skill_prompt),
272 None => skill_prompt,
273 };
274 ab = ab.system_prompt(new_prompt);
275 }
276 }
277
278 let detail_tool = SkillDetailTool::new(skill_refs, self.skill_detail_tool_name);
279 ab = ab.register_tool(detail_tool);
280
281 ab.build()
282 }
283}