Skip to main content

adk_ui/
toolset.rs

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