gemini_adk_fluent_rs/compose/
tools.rs

1//! T — Tool composition.
2//!
3//! Compose tools in any order with `|`.
4
5use std::future::Future;
6use std::pin::Pin;
7use std::sync::Arc;
8
9use gemini_adk_rs::text::TextAgent;
10use gemini_adk_rs::tool::{PolicyTool, SimpleTool, ToolFunction, ToolPolicy};
11use gemini_genai_rs::prelude::{FunctionDeclaration, Tool};
12
13/// A tool composite — one or more tool entries.
14///
15/// Built from the `T` namespace and composed with `|`. Any single
16/// [`ToolFunction`] (a `SimpleTool`, a `TypedTool`, the value a `#[tool]`
17/// function returns, or an `Arc<dyn ToolFunction>`) converts into a
18/// one-entry composite via `From`, so `.tools(get_weather())` works without
19/// the namespace.
20#[derive(Clone)]
21#[non_exhaustive]
22pub struct ToolComposite {
23    /// The tool entries in this composite.
24    pub entries: Vec<ToolCompositeEntry>,
25}
26
27/// Async transformer applied to a tool result value.
28pub type TransformFn = Arc<
29    dyn Fn(serde_json::Value) -> Pin<Box<dyn Future<Output = serde_json::Value> + Send>>
30        + Send
31        + Sync,
32>;
33
34/// An entry in a tool composite.
35#[derive(Clone)]
36pub enum ToolCompositeEntry {
37    /// A runtime tool function.
38    Function(Arc<dyn ToolFunction>),
39    /// A built-in Gemini tool declaration.
40    BuiltIn(Tool),
41    /// A text agent wrapped as a tool.
42    Agent {
43        /// Tool name exposed to the model.
44        name: String,
45        /// Tool description exposed to the model.
46        description: String,
47        /// The text agent to invoke.
48        agent: Arc<dyn TextAgent>,
49    },
50    /// An MCP (Model Context Protocol) toolset connection.
51    Mcp {
52        /// Connection params (e.g. URL or command string).
53        params: String,
54    },
55    /// A mock tool that returns a fixed response (useful for testing).
56    Mock {
57        /// Tool name.
58        name: String,
59        /// Tool description.
60        description: String,
61        /// Fixed response to return.
62        response: serde_json::Value,
63    },
64    /// A schema-defined tool (placeholder/marker).
65    Schema {
66        /// Tool name.
67        name: String,
68        /// JSON Schema defining the tool's parameters.
69        schema: serde_json::Value,
70    },
71    /// A tool wrapped with a result transformer.
72    Transform {
73        /// The inner tool entry.
74        inner: Box<ToolCompositeEntry>,
75        /// Transformer function applied to the tool result.
76        transformer: TransformFn,
77    },
78}
79
80impl ToolComposite {
81    /// Create a composite containing a single runtime tool function.
82    pub fn from_function(f: Arc<dyn ToolFunction>) -> Self {
83        Self {
84            entries: vec![ToolCompositeEntry::Function(f)],
85        }
86    }
87
88    /// Create a composite containing a single built-in tool declaration.
89    pub fn from_built_in(tool: Tool) -> Self {
90        Self {
91            entries: vec![ToolCompositeEntry::BuiltIn(tool)],
92        }
93    }
94
95    /// Number of tool entries.
96    pub fn len(&self) -> usize {
97        self.entries.len()
98    }
99
100    /// Whether empty.
101    pub fn is_empty(&self) -> bool {
102        self.entries.is_empty()
103    }
104
105    /// Apply a per-tool [`ToolPolicy`] transform to every function entry.
106    ///
107    /// Each [`ToolCompositeEntry::Function`] is wrapped in a [`PolicyTool`]
108    /// carrying the policy. Successive modifiers nest (e.g. `T::cached(T::timeout(..))`
109    /// applies both timeout and cache), since a `PolicyTool` is itself a
110    /// [`ToolFunction`]. Other entry kinds are left untouched.
111    fn map_function_policy(
112        mut self,
113        f: impl Fn(ToolPolicy) -> ToolPolicy + Send + Sync + 'static,
114    ) -> Self {
115        self.entries = self
116            .entries
117            .into_iter()
118            .map(|entry| match entry {
119                ToolCompositeEntry::Function(func) => {
120                    let policy = f(ToolPolicy::new());
121                    ToolCompositeEntry::Function(PolicyTool::wrap(func, policy))
122                }
123                other => other,
124            })
125            .collect();
126        self
127    }
128}
129
130/// A single tool is a one-entry composite, so `.tools(my_tool)` and
131/// `.tool(my_tool)` accept the same values.
132impl<F: ToolFunction + 'static> From<F> for ToolComposite {
133    fn from(f: F) -> Self {
134        Self::from_function(Arc::new(f))
135    }
136}
137
138/// Compose two tool composites with `|`.
139impl std::ops::BitOr for ToolComposite {
140    type Output = ToolComposite;
141
142    fn bitor(mut self, rhs: ToolComposite) -> Self::Output {
143        self.entries.extend(rhs.entries);
144        self
145    }
146}
147
148/// The `T` namespace — static factory methods for tool composition.
149pub struct T;
150
151impl T {
152    /// Register a function tool.
153    pub fn function(f: Arc<dyn ToolFunction>) -> ToolComposite {
154        ToolComposite::from_function(f)
155    }
156
157    /// Add Google Search built-in tool.
158    pub fn google_search() -> ToolComposite {
159        ToolComposite::from_built_in(Tool::google_search())
160    }
161
162    /// Add URL context built-in tool.
163    pub fn url_context() -> ToolComposite {
164        ToolComposite::from_built_in(Tool::url_context())
165    }
166
167    /// Add code execution built-in tool.
168    pub fn code_execution() -> ToolComposite {
169        ToolComposite::from_built_in(Tool::code_execution())
170    }
171
172    /// Create a simple tool from a name, description, and async closure.
173    pub fn simple<F, Fut>(
174        name: impl Into<String>,
175        description: impl Into<String>,
176        f: F,
177    ) -> ToolComposite
178    where
179        F: Fn(serde_json::Value) -> Fut + Send + Sync + 'static,
180        Fut: Future<Output = Result<serde_json::Value, gemini_adk_rs::ToolError>> + Send + 'static,
181    {
182        let tool = SimpleTool::new(name, description, None, f);
183        ToolComposite::from_function(Arc::new(tool))
184    }
185
186    /// Alias for [`simple`](Self::simple) — matches upstream Python `T.fn()`.
187    ///
188    /// Named `fn_tool` because `fn` is a reserved keyword in Rust.
189    pub fn fn_tool<F, Fut>(
190        name: impl Into<String>,
191        description: impl Into<String>,
192        f: F,
193    ) -> ToolComposite
194    where
195        F: Fn(serde_json::Value) -> Fut + Send + Sync + 'static,
196        Fut: Future<Output = Result<serde_json::Value, gemini_adk_rs::ToolError>> + Send + 'static,
197    {
198        Self::simple(name, description, f)
199    }
200
201    /// Require user confirmation before each function tool in the composite runs.
202    ///
203    /// The confirmation flag is recorded on the tool's [`ToolPolicy`] and surfaced
204    /// to the runtime via [`PolicyTool::requires_confirmation`] — it is never
205    /// silently dropped. The `message` becomes the confirmation hint. Built-in and
206    /// placeholder entries are left unchanged.
207    pub fn confirm(tool: ToolComposite, message: &str) -> ToolComposite {
208        let msg = if message.is_empty() {
209            None
210        } else {
211            Some(message.to_string())
212        };
213        tool.map_function_policy(move |p| p.with_confirm(msg.clone()))
214    }
215
216    /// Bound each function tool in the composite by a timeout.
217    ///
218    /// At dispatch the tool's future is raced against the duration; on elapse the
219    /// call returns [`ToolError::Timeout`](gemini_adk_rs::ToolError::Timeout).
220    /// Built-in and placeholder entries are left unchanged.
221    pub fn timeout(tool: ToolComposite, duration: std::time::Duration) -> ToolComposite {
222        tool.map_function_policy(move |p| p.with_timeout(duration))
223    }
224
225    /// Memoize each function tool's successful results.
226    ///
227    /// Results are cached by `(tool name, canonical-JSON args)`; repeat calls with
228    /// identical arguments return the cached value without re-invoking the tool.
229    /// Errors are not cached. Built-in/placeholder entries are left unchanged.
230    pub fn cached(tool: ToolComposite) -> ToolComposite {
231        tool.map_function_policy(gemini_adk_rs::tool::ToolPolicy::with_cache)
232    }
233
234    /// Combine multiple tool functions into a single composite.
235    pub fn toolset(tools: Vec<Arc<dyn ToolFunction>>) -> ToolComposite {
236        ToolComposite {
237            entries: tools
238                .into_iter()
239                .map(ToolCompositeEntry::Function)
240                .collect(),
241        }
242    }
243
244    /// Wrap a [`TextAgent`] as a tool (shorthand for creating an agent tool entry).
245    ///
246    /// When invoked, the agent runs via `BaseLlm::generate()` and returns its
247    /// text output as the tool result. State is shared with the parent session.
248    pub fn agent(
249        name: impl Into<String>,
250        description: impl Into<String>,
251        agent: impl TextAgent + 'static,
252    ) -> ToolComposite {
253        ToolComposite {
254            entries: vec![ToolCompositeEntry::Agent {
255                name: name.into(),
256                description: description.into(),
257                agent: Arc::new(agent),
258            }],
259        }
260    }
261
262    /// Create an MCP (Model Context Protocol) toolset entry.
263    ///
264    /// `params` is the connection string (e.g. a URL or command) used to
265    /// establish the MCP session at runtime.
266    pub fn mcp(params: impl Into<String>) -> ToolComposite {
267        ToolComposite {
268            entries: vec![ToolCompositeEntry::Mcp {
269                params: params.into(),
270            }],
271        }
272    }
273
274    /// Create a mock tool that returns a fixed response.
275    ///
276    /// Useful for testing and prototyping without real tool implementations.
277    pub fn mock(
278        name: impl Into<String>,
279        description: impl Into<String>,
280        response: serde_json::Value,
281    ) -> ToolComposite {
282        ToolComposite {
283            entries: vec![ToolCompositeEntry::Mock {
284                name: name.into(),
285                description: description.into(),
286                response,
287            }],
288        }
289    }
290
291    /// Create a schema-defined tool (placeholder/marker).
292    ///
293    /// The tool's parameters are defined by the given JSON Schema value.
294    pub fn schema(name: impl Into<String>, schema: serde_json::Value) -> ToolComposite {
295        ToolComposite {
296            entries: vec![ToolCompositeEntry::Schema {
297                name: name.into(),
298                schema,
299            }],
300        }
301    }
302
303    /// Wrap each tool entry in a composite with a result transformer.
304    ///
305    /// The transformer function is applied to the tool's output value before
306    /// it is returned to the model.
307    pub fn transform<F, Fut>(tool: ToolComposite, f: F) -> ToolComposite
308    where
309        F: Fn(serde_json::Value) -> Fut + Send + Sync + 'static,
310        Fut: Future<Output = serde_json::Value> + Send + 'static,
311    {
312        let f: TransformFn = Arc::new(
313            move |v: serde_json::Value| -> Pin<Box<dyn Future<Output = serde_json::Value> + Send>> {
314                Box::pin(f(v))
315            },
316        );
317        ToolComposite {
318            entries: tool
319                .entries
320                .into_iter()
321                .map(|entry| ToolCompositeEntry::Transform {
322                    inner: Box::new(entry),
323                    transformer: Arc::clone(&f),
324                })
325                .collect(),
326        }
327    }
328}
329
330// ── Resolution ────────────────────────────────────────────────────────────
331
332/// A tool entry that needs asynchronous I/O (network or subprocess) to resolve,
333/// and is therefore resolved at connect time rather than when the composite is
334/// built. See [`crate::live::Live`] connection methods. A text
335/// [`AgentBuilder`](crate::builder::AgentBuilder) rejects these at `build`.
336#[derive(Clone, Debug)]
337pub enum DeferredTool {
338    /// MCP server connection — a stdio command line or an SSE/HTTP URL.
339    Mcp {
340        /// Connection string: an `http(s)://` URL (SSE) or a command line (stdio).
341        params: String,
342    },
343}
344
345/// The concrete outcome of classifying a single [`ToolCompositeEntry`].
346///
347/// This is the *single* exhaustive mapping from the composable tool algebra to
348/// the runtime; both [`crate::builder::AgentBuilder`] and [`crate::live::Live`]
349/// resolve through it, so no entry can be silently dropped.
350pub(crate) enum ToolResolution {
351    /// A runtime-executable tool function (register with a dispatcher).
352    Runtime(Arc<dyn ToolFunction>),
353    /// A built-in / declaration-only Gemini tool (add to the session config).
354    BuiltIn(Tool),
355    /// A text agent to expose as a tool (needs a shared session `State`).
356    Agent {
357        /// Tool name exposed to the model.
358        name: String,
359        /// Tool description exposed to the model.
360        description: String,
361        /// The text agent to invoke.
362        agent: Arc<dyn TextAgent>,
363    },
364    /// A tool that can only be resolved with async I/O at connect time.
365    Deferred(DeferredTool),
366}
367
368impl ToolCompositeEntry {
369    #[cfg(test)]
370    fn classify_name(self) -> String {
371        match self.classify() {
372            ToolResolution::Runtime(f) => f.name().to_string(),
373            _ => String::new(),
374        }
375    }
376
377    /// Classify this entry into its concrete [`ToolResolution`]. Exhaustive by
378    /// construction — adding a variant forces every consumer to handle it.
379    pub(crate) fn classify(self) -> ToolResolution {
380        match self {
381            ToolCompositeEntry::Function(f) => ToolResolution::Runtime(f),
382            ToolCompositeEntry::BuiltIn(t) => ToolResolution::BuiltIn(t),
383            ToolCompositeEntry::Agent {
384                name,
385                description,
386                agent,
387            } => ToolResolution::Agent {
388                name,
389                description,
390                agent,
391            },
392            ToolCompositeEntry::Mock {
393                name,
394                description,
395                response,
396            } => ToolResolution::Runtime(Arc::new(SimpleTool::new(
397                name,
398                description,
399                None,
400                move |_args| {
401                    let r = response.clone();
402                    async move { Ok(r) }
403                },
404            ))),
405            ToolCompositeEntry::Transform { inner, transformer } => match inner.classify() {
406                ToolResolution::Runtime(f) => ToolResolution::Runtime(Arc::new(TransformTool {
407                    inner: f,
408                    transformer,
409                })),
410                // A transformer only applies to a runtime function; for any other
411                // inner kind the transform is a no-op and the inner resolution
412                // passes through unchanged.
413                other => other,
414            },
415            ToolCompositeEntry::Schema { name, schema } => {
416                // A declaration-only tool: the model is told the function exists
417                // and the application services the call (e.g. via on_tool_call).
418                ToolResolution::BuiltIn(Tool::functions(vec![FunctionDeclaration {
419                    name,
420                    description: String::new(),
421                    parameters: Some(schema),
422                    behavior: None,
423                }]))
424            }
425            ToolCompositeEntry::Mcp { params } => {
426                ToolResolution::Deferred(DeferredTool::Mcp { params })
427            }
428        }
429    }
430}
431
432/// A [`ToolFunction`] that applies an async transformer to another tool's result.
433struct TransformTool {
434    inner: Arc<dyn ToolFunction>,
435    #[allow(clippy::type_complexity)]
436    transformer: Arc<
437        dyn Fn(serde_json::Value) -> Pin<Box<dyn Future<Output = serde_json::Value> + Send>>
438            + Send
439            + Sync,
440    >,
441}
442
443#[async_trait::async_trait]
444impl ToolFunction for TransformTool {
445    fn name(&self) -> &str {
446        self.inner.name()
447    }
448
449    fn description(&self) -> &str {
450        self.inner.description()
451    }
452
453    fn parameters(&self) -> Option<serde_json::Value> {
454        self.inner.parameters()
455    }
456
457    async fn call(
458        &self,
459        args: serde_json::Value,
460    ) -> Result<serde_json::Value, gemini_adk_rs::error::ToolError> {
461        let result = self.inner.call(args).await?;
462        Ok((self.transformer)(result).await)
463    }
464}
465
466#[cfg(test)]
467mod tests {
468    use super::*;
469
470    /// Classify the single entry of a one-element composite.
471    fn classify_one(c: ToolComposite) -> ToolResolution {
472        c.entries.into_iter().next().unwrap().classify()
473    }
474
475    #[test]
476    fn classify_maps_every_variant() {
477        // Synchronous, runtime-executable.
478        assert!(matches!(
479            classify_one(T::mock("m", "d", serde_json::json!({"ok": true}))),
480            ToolResolution::Runtime(_)
481        ));
482        assert!(matches!(
483            classify_one(T::simple("s", "d", |a| async move { Ok(a) })),
484            ToolResolution::Runtime(_)
485        ));
486        // Built-in / declaration-only.
487        assert!(matches!(
488            classify_one(T::google_search()),
489            ToolResolution::BuiltIn(_)
490        ));
491        assert!(matches!(
492            classify_one(T::schema("s", serde_json::json!({"type": "object"}))),
493            ToolResolution::BuiltIn(_)
494        ));
495        // Async, connect-time.
496        assert!(matches!(
497            classify_one(T::mcp("node ./server.js")),
498            ToolResolution::Deferred(DeferredTool::Mcp { .. })
499        ));
500    }
501
502    #[test]
503    fn a_single_tool_function_converts_into_a_composite() {
504        let composite: ToolComposite =
505            SimpleTool::new("one", "one", None, |_| async { Ok(serde_json::json!(1)) }).into();
506        assert_eq!(composite.len(), 1);
507        let arc: Arc<dyn ToolFunction> = Arc::new(SimpleTool::new("two", "two", None, |_| async {
508            Ok(serde_json::json!(2))
509        }));
510        let composite: ToolComposite = arc.into();
511        assert_eq!(composite.len(), 1);
512        assert_eq!(composite.entries[0].clone().classify_name(), "two");
513    }
514
515    #[tokio::test]
516    async fn mock_resolves_to_callable_runtime_tool() {
517        let resolution = classify_one(T::mock(
518            "weather",
519            "Mock weather",
520            serde_json::json!({"temp": 22}),
521        ));
522        let ToolResolution::Runtime(tool) = resolution else {
523            panic!("mock should resolve to a runtime tool");
524        };
525        assert_eq!(tool.name(), "weather");
526        let out = tool.call(serde_json::json!({})).await.unwrap();
527        assert_eq!(out, serde_json::json!({"temp": 22}));
528    }
529
530    #[tokio::test]
531    async fn transform_wraps_inner_runtime_result() {
532        let composite = T::transform(
533            T::mock("base", "d", serde_json::json!({"n": 1})),
534            |mut v| async move {
535                v["doubled"] = serde_json::json!(true);
536                v
537            },
538        );
539        let ToolResolution::Runtime(tool) = classify_one(composite) else {
540            panic!("transform over a mock should resolve to a runtime tool");
541        };
542        assert_eq!(tool.name(), "base");
543        let out = tool.call(serde_json::json!({})).await.unwrap();
544        assert_eq!(out, serde_json::json!({"n": 1, "doubled": true}));
545    }
546
547    #[test]
548    fn google_search_creates_composite() {
549        let t = T::google_search();
550        assert_eq!(t.len(), 1);
551    }
552
553    #[test]
554    fn url_context_creates_composite() {
555        let t = T::url_context();
556        assert_eq!(t.len(), 1);
557    }
558
559    #[test]
560    fn code_execution_creates_composite() {
561        let t = T::code_execution();
562        assert_eq!(t.len(), 1);
563    }
564
565    #[test]
566    fn compose_with_bitor() {
567        let t = T::google_search() | T::url_context() | T::code_execution();
568        assert_eq!(t.len(), 3);
569    }
570
571    #[test]
572    fn simple_creates_tool() {
573        let t = T::simple("greet", "Greets the user", |_args| async {
574            Ok(serde_json::json!({"message": "hello"}))
575        });
576        assert_eq!(t.len(), 1);
577        match &t.entries[0] {
578            ToolCompositeEntry::Function(f) => assert_eq!(f.name(), "greet"),
579            _ => panic!("expected Function entry"),
580        }
581    }
582
583    #[tokio::test]
584    async fn timeout_modifier_enforces_timeout() {
585        use gemini_adk_rs::ToolError;
586        use std::time::Duration;
587
588        let t = T::timeout(
589            T::simple("slow", "slow tool", |_| async move {
590                tokio::time::sleep(Duration::from_secs(3600)).await;
591                Ok(serde_json::json!({"ok": true}))
592            }),
593            Duration::from_millis(50),
594        );
595        match &t.entries[0] {
596            ToolCompositeEntry::Function(f) => match f.call(serde_json::json!({})).await {
597                Err(ToolError::Timeout(d)) => assert_eq!(d, Duration::from_millis(50)),
598                other => panic!("expected Timeout, got {other:?}"),
599            },
600            _ => panic!("expected Function entry"),
601        }
602    }
603
604    #[tokio::test]
605    async fn cached_modifier_memoizes_results() {
606        use std::sync::atomic::{AtomicU32, Ordering};
607
608        let counter = Arc::new(AtomicU32::new(0));
609        let c = counter.clone();
610        let t = T::cached(T::simple("count", "counts calls", move |_| {
611            let c = c.clone();
612            async move {
613                let n = c.fetch_add(1, Ordering::SeqCst) + 1;
614                Ok(serde_json::json!({"n": n}))
615            }
616        }));
617        match &t.entries[0] {
618            ToolCompositeEntry::Function(f) => {
619                let first = f.call(serde_json::json!({"x": 1})).await.unwrap();
620                let second = f.call(serde_json::json!({"x": 1})).await.unwrap();
621                assert_eq!(first, second);
622                assert_eq!(first["n"], 1);
623                assert_eq!(counter.load(Ordering::SeqCst), 1);
624            }
625            _ => panic!("expected Function entry"),
626        }
627    }
628
629    #[test]
630    fn confirm_modifier_wraps_function() {
631        // confirm() wraps the function (preserving its name) so the policy flag
632        // travels to the runtime rather than being silently dropped.
633        let t = T::confirm(
634            T::simple("danger", "dangerous", |_| async move {
635                Ok(serde_json::json!({}))
636            }),
637            "are you sure?",
638        );
639        match &t.entries[0] {
640            ToolCompositeEntry::Function(f) => assert_eq!(f.name(), "danger"),
641            _ => panic!("expected Function entry"),
642        }
643    }
644
645    #[test]
646    fn toolset_combines_functions() {
647        let tool_a: Arc<dyn ToolFunction> =
648            Arc::new(SimpleTool::new("a", "tool a", None, |_| async {
649                Ok(serde_json::json!(null))
650            }));
651        let tool_b: Arc<dyn ToolFunction> =
652            Arc::new(SimpleTool::new("b", "tool b", None, |_| async {
653                Ok(serde_json::json!(null))
654            }));
655        let t = T::toolset(vec![tool_a, tool_b]);
656        assert_eq!(t.len(), 2);
657    }
658}
gemini_adk_fluent_rs/compose/tools.rs

gemini_adk_fluent_rs/compose/
tools.rs
