Skip to main content

adk_ui/
toolset.rs

1use crate::compat::{ReadonlyContext, Result, Tool, Toolset};
2use crate::tools::*;
3use async_trait::async_trait;
4use std::sync::Arc;
5
6/// A toolset containing all UI rendering tools.
7///
8/// Use this to easily add UI capabilities to an agent:
9///
10/// ```rust,ignore
11/// use adk_ui::UiToolset;
12/// use adk_agent::LlmAgentBuilder;
13///
14/// let tools = UiToolset::all_tools();
15/// let mut builder = LlmAgentBuilder::new("assistant");
16/// for tool in tools {
17///     builder = builder.tool(tool);
18/// }
19/// let agent = builder.build()?;
20/// ```
21pub struct UiToolset {
22    include_screen: bool,
23    include_app: bool,
24    include_page: bool,
25    include_kit: bool,
26    include_form: bool,
27    include_card: bool,
28    include_alert: bool,
29    include_confirm: bool,
30    include_table: bool,
31    include_chart: bool,
32    include_layout: bool,
33    include_progress: bool,
34    include_modal: bool,
35    include_toast: bool,
36}
37
38impl UiToolset {
39    /// Create a new UiToolset with all tools enabled
40    pub fn new() -> Self {
41        Self {
42            include_screen: true,
43            include_app: true,
44            include_page: true,
45            include_kit: true,
46            include_form: true,
47            include_card: true,
48            include_alert: true,
49            include_confirm: true,
50            include_table: true,
51            include_chart: true,
52            include_layout: true,
53            include_progress: true,
54            include_modal: true,
55            include_toast: true,
56        }
57    }
58
59    /// Create a toolset with only form rendering
60    pub fn forms_only() -> Self {
61        Self {
62            include_screen: false,
63            include_app: false,
64            include_page: false,
65            include_kit: false,
66            include_form: true,
67            include_card: false,
68            include_alert: false,
69            include_confirm: false,
70            include_table: false,
71            include_chart: false,
72            include_layout: false,
73            include_progress: false,
74            include_modal: false,
75            include_toast: false,
76        }
77    }
78
79    /// Disable form rendering
80    pub fn without_form(mut self) -> Self {
81        self.include_form = false;
82        self
83    }
84
85    /// Disable screen rendering
86    pub fn without_screen(mut self) -> Self {
87        self.include_screen = false;
88        self
89    }
90
91    /// Disable full application rendering
92    pub fn without_app(mut self) -> Self {
93        self.include_app = false;
94        self
95    }
96
97    /// Disable page rendering
98    pub fn without_page(mut self) -> Self {
99        self.include_page = false;
100        self
101    }
102
103    /// Disable kit rendering
104    pub fn without_kit(mut self) -> Self {
105        self.include_kit = false;
106        self
107    }
108
109    /// Disable card rendering
110    pub fn without_card(mut self) -> Self {
111        self.include_card = false;
112        self
113    }
114
115    /// Disable alert rendering
116    pub fn without_alert(mut self) -> Self {
117        self.include_alert = false;
118        self
119    }
120
121    /// Disable confirm rendering
122    pub fn without_confirm(mut self) -> Self {
123        self.include_confirm = false;
124        self
125    }
126
127    /// Disable table rendering
128    pub fn without_table(mut self) -> Self {
129        self.include_table = false;
130        self
131    }
132
133    /// Disable chart rendering
134    pub fn without_chart(mut self) -> Self {
135        self.include_chart = false;
136        self
137    }
138
139    /// Disable layout rendering
140    pub fn without_layout(mut self) -> Self {
141        self.include_layout = false;
142        self
143    }
144
145    /// Disable progress rendering
146    pub fn without_progress(mut self) -> Self {
147        self.include_progress = false;
148        self
149    }
150
151    /// Disable modal rendering
152    pub fn without_modal(mut self) -> Self {
153        self.include_modal = false;
154        self
155    }
156
157    /// Disable toast rendering
158    pub fn without_toast(mut self) -> Self {
159        self.include_toast = false;
160        self
161    }
162
163    /// Get all tools as a Vec for use with LlmAgentBuilder
164    pub fn all_tools() -> Vec<Arc<dyn Tool>> {
165        vec![
166            Arc::new(RenderScreenTool::new()) as Arc<dyn Tool>,
167            Arc::new(RenderAppTool::new()),
168            Arc::new(RenderPageTool::new()),
169            Arc::new(RenderKitTool::new()),
170            Arc::new(RenderFormTool::new()) as Arc<dyn Tool>,
171            Arc::new(RenderCardTool::new()),
172            Arc::new(RenderAlertTool::new()),
173            Arc::new(RenderConfirmTool::new()),
174            Arc::new(RenderTableTool::new()),
175            Arc::new(RenderChartTool::new()),
176            Arc::new(RenderLayoutTool::new()),
177            Arc::new(RenderProgressTool::new()),
178            Arc::new(RenderModalTool::new()),
179            Arc::new(RenderToastTool::new()),
180        ]
181    }
182}
183
184#[cfg(feature = "awp")]
185impl UiToolset {
186    /// Export enabled tools as AWP CapabilityEntry values.
187    ///
188    /// Uses the per-tool include flags to determine which tools to export.
189    /// Does not require a `ReadonlyContext` or async context.
190    pub fn to_capability_entries(&self) -> Vec<awp_types::CapabilityEntry> {
191        let mut entries = Vec::new();
192
193        macro_rules! maybe_export {
194            ($flag:expr, $tool:expr) => {
195                if $flag {
196                    let tool = $tool;
197                    entries.push(awp_types::CapabilityEntry {
198                        name: tool.name().to_string(),
199                        description: tool.description().to_string(),
200                        endpoint: format!("/tools/{}", tool.name()),
201                        method: "POST".to_string(),
202                        input_schema: tool.parameters_schema().map(|v| v.to_string()),
203                        output_schema: None,
204                    });
205                }
206            };
207        }
208
209        maybe_export!(self.include_screen, RenderScreenTool::new());
210        maybe_export!(self.include_app, RenderAppTool::new());
211        maybe_export!(self.include_page, RenderPageTool::new());
212        maybe_export!(self.include_kit, RenderKitTool::new());
213        maybe_export!(self.include_form, RenderFormTool::new());
214        maybe_export!(self.include_card, RenderCardTool::new());
215        maybe_export!(self.include_alert, RenderAlertTool::new());
216        maybe_export!(self.include_confirm, RenderConfirmTool::new());
217        maybe_export!(self.include_table, RenderTableTool::new());
218        maybe_export!(self.include_chart, RenderChartTool::new());
219        maybe_export!(self.include_layout, RenderLayoutTool::new());
220        maybe_export!(self.include_progress, RenderProgressTool::new());
221        maybe_export!(self.include_modal, RenderModalTool::new());
222        maybe_export!(self.include_toast, RenderToastTool::new());
223
224        entries
225    }
226}
227
228impl Default for UiToolset {
229    fn default() -> Self {
230        Self::new()
231    }
232}
233
234#[async_trait]
235impl Toolset for UiToolset {
236    fn name(&self) -> &str {
237        "ui"
238    }
239
240    async fn tools(&self, _ctx: Arc<dyn ReadonlyContext>) -> Result<Vec<Arc<dyn Tool>>> {
241        let mut tools: Vec<Arc<dyn Tool>> = Vec::new();
242
243        if self.include_screen {
244            tools.push(Arc::new(RenderScreenTool::new()));
245        }
246        if self.include_app {
247            tools.push(Arc::new(RenderAppTool::new()));
248        }
249        if self.include_page {
250            tools.push(Arc::new(RenderPageTool::new()));
251        }
252        if self.include_kit {
253            tools.push(Arc::new(RenderKitTool::new()));
254        }
255        if self.include_form {
256            tools.push(Arc::new(RenderFormTool::new()));
257        }
258        if self.include_card {
259            tools.push(Arc::new(RenderCardTool::new()));
260        }
261        if self.include_alert {
262            tools.push(Arc::new(RenderAlertTool::new()));
263        }
264        if self.include_confirm {
265            tools.push(Arc::new(RenderConfirmTool::new()));
266        }
267        if self.include_table {
268            tools.push(Arc::new(RenderTableTool::new()));
269        }
270        if self.include_chart {
271            tools.push(Arc::new(RenderChartTool::new()));
272        }
273        if self.include_layout {
274            tools.push(Arc::new(RenderLayoutTool::new()));
275        }
276        if self.include_progress {
277            tools.push(Arc::new(RenderProgressTool::new()));
278        }
279        if self.include_modal {
280            tools.push(Arc::new(RenderModalTool::new()));
281        }
282        if self.include_toast {
283            tools.push(Arc::new(RenderToastTool::new()));
284        }
285
286        Ok(tools)
287    }
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293
294    #[test]
295    fn test_all_tools_returns_14_tools() {
296        let tools = UiToolset::all_tools();
297        assert_eq!(tools.len(), 14);
298
299        let names: Vec<&str> = tools.iter().map(|t| t.name()).collect();
300        assert!(names.contains(&"render_screen"));
301        assert!(names.contains(&"render_app"));
302        assert!(names.contains(&"render_page"));
303        assert!(names.contains(&"render_kit"));
304        assert!(names.contains(&"render_form"));
305        assert!(names.contains(&"render_card"));
306        assert!(names.contains(&"render_alert"));
307        assert!(names.contains(&"render_confirm"));
308        assert!(names.contains(&"render_table"));
309        assert!(names.contains(&"render_chart"));
310        assert!(names.contains(&"render_layout"));
311        assert!(names.contains(&"render_progress"));
312        assert!(names.contains(&"render_modal"));
313        assert!(names.contains(&"render_toast"));
314    }
315
316    #[test]
317    fn test_forms_only() {
318        let toolset = UiToolset::forms_only();
319        assert!(!toolset.include_screen);
320        assert!(!toolset.include_app);
321        assert!(!toolset.include_page);
322        assert!(!toolset.include_kit);
323        assert!(toolset.include_form);
324        assert!(!toolset.include_card);
325        assert!(!toolset.include_alert);
326        assert!(!toolset.include_table);
327    }
328
329    #[test]
330    fn test_without_methods() {
331        let toolset = UiToolset::new()
332            .without_chart()
333            .without_table()
334            .without_progress();
335
336        assert!(toolset.include_form);
337        assert!(toolset.include_card);
338        assert!(!toolset.include_chart);
339        assert!(!toolset.include_table);
340        assert!(!toolset.include_progress);
341    }
342
343    #[test]
344    fn test_toolset_name() {
345        let toolset = UiToolset::new();
346        assert_eq!(toolset.name(), "ui");
347    }
348
349    #[test]
350    fn test_default_is_new() {
351        let default = UiToolset::default();
352        let new = UiToolset::new();
353        assert_eq!(default.include_screen, new.include_screen);
354        assert_eq!(default.include_page, new.include_page);
355        assert_eq!(default.include_kit, new.include_kit);
356        assert_eq!(default.include_form, new.include_form);
357        assert_eq!(default.include_card, new.include_card);
358        assert_eq!(default.include_chart, new.include_chart);
359    }
360}