1use 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, TypedTool};
11use gemini_genai_rs::prelude::{FunctionDeclaration, Tool};
12
13#[derive(Clone)]
21#[non_exhaustive]
22pub struct ToolComposite {
23 pub entries: Vec<ToolCompositeEntry>,
25}
26
27pub 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#[derive(Clone)]
36pub enum ToolCompositeEntry {
37 Function(Arc<dyn ToolFunction>),
39 BuiltIn(Tool),
41 Agent {
43 name: String,
45 description: String,
47 agent: Arc<dyn TextAgent>,
49 },
50 Mcp {
52 params: String,
54 },
55 Mock {
57 name: String,
59 description: String,
61 response: serde_json::Value,
63 },
64 Schema {
66 name: String,
68 schema: serde_json::Value,
70 },
71 Transform {
73 inner: Box<ToolCompositeEntry>,
75 transformer: TransformFn,
77 },
78}
79
80impl ToolComposite {
81 pub fn from_function(f: Arc<dyn ToolFunction>) -> Self {
83 Self {
84 entries: vec![ToolCompositeEntry::Function(f)],
85 }
86 }
87
88 pub fn from_built_in(tool: Tool) -> Self {
90 Self {
91 entries: vec![ToolCompositeEntry::BuiltIn(tool)],
92 }
93 }
94
95 pub fn len(&self) -> usize {
97 self.entries.len()
98 }
99
100 pub fn is_empty(&self) -> bool {
102 self.entries.is_empty()
103 }
104
105 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
130impl<F: ToolFunction + 'static> From<F> for ToolComposite {
133 fn from(f: F) -> Self {
134 Self::from_function(Arc::new(f))
135 }
136}
137
138impl std::ops::Add for ToolComposite {
140 type Output = ToolComposite;
141
142 fn add(mut self, rhs: ToolComposite) -> Self::Output {
143 self.entries.extend(rhs.entries);
144 self
145 }
146}
147
148pub struct T;
150
151impl T {
152 pub fn function(f: Arc<dyn ToolFunction>) -> ToolComposite {
154 ToolComposite::from_function(f)
155 }
156
157 pub fn google_search() -> ToolComposite {
159 ToolComposite::from_built_in(Tool::google_search())
160 }
161
162 pub fn url_context() -> ToolComposite {
164 ToolComposite::from_built_in(Tool::url_context())
165 }
166
167 pub fn code_execution() -> ToolComposite {
169 ToolComposite::from_built_in(Tool::code_execution())
170 }
171
172 pub fn simple<F, Fut>(
180 name: impl Into<String>,
181 description: impl Into<String>,
182 f: F,
183 ) -> ToolComposite
184 where
185 F: Fn(serde_json::Value) -> Fut + Send + Sync + 'static,
186 Fut: Future<Output = Result<serde_json::Value, gemini_adk_rs::ToolError>> + Send + 'static,
187 {
188 let tool = SimpleTool::new(name, description, None, f);
189 ToolComposite::from_function(Arc::new(tool))
190 }
191
192 pub fn contextual<F, Fut>(
205 name: impl Into<String>,
206 description: impl Into<String>,
207 f: F,
208 ) -> ToolComposite
209 where
210 F: Fn(serde_json::Value, gemini_adk_rs::tool::ToolContext) -> Fut + Send + Sync + 'static,
211 Fut: Future<Output = Result<serde_json::Value, gemini_adk_rs::ToolError>> + Send + 'static,
212 {
213 let tool = gemini_adk_rs::tool::ContextTool::new(name, description, None, f);
214 ToolComposite::from_function(Arc::new(tool))
215 }
216
217 pub fn typed<A, F, Fut>(
245 name: impl Into<String>,
246 description: impl Into<String>,
247 f: F,
248 ) -> ToolComposite
249 where
250 A: serde::de::DeserializeOwned + schemars::JsonSchema + Send + Sync + 'static,
251 F: Fn(A) -> Fut + Send + Sync + 'static,
252 Fut: Future<Output = Result<serde_json::Value, gemini_adk_rs::ToolError>> + Send + 'static,
253 {
254 ToolComposite::from_function(Arc::new(TypedTool::new::<F, Fut>(name, description, f)))
255 }
256
257 #[deprecated(
259 since = "2.1.0",
260 note = "use `T::simple` (no parameters), `T::typed` or `#[tool]`"
261 )]
262 pub fn fn_tool<F, Fut>(
265 name: impl Into<String>,
266 description: impl Into<String>,
267 f: F,
268 ) -> ToolComposite
269 where
270 F: Fn(serde_json::Value) -> Fut + Send + Sync + 'static,
271 Fut: Future<Output = Result<serde_json::Value, gemini_adk_rs::ToolError>> + Send + 'static,
272 {
273 Self::simple(name, description, f)
274 }
275
276 pub fn confirm(tool: impl Into<ToolComposite>, message: &str) -> ToolComposite {
283 let tool = tool.into();
284 let msg = if message.is_empty() {
285 None
286 } else {
287 Some(message.to_string())
288 };
289 tool.map_function_policy(move |p| p.with_confirm(msg.clone()))
290 }
291
292 pub fn timeout(tool: impl Into<ToolComposite>, duration: std::time::Duration) -> ToolComposite {
298 tool.into()
299 .map_function_policy(move |p| p.with_timeout(duration))
300 }
301
302 pub fn cached(tool: impl Into<ToolComposite>) -> ToolComposite {
308 tool.into()
309 .map_function_policy(gemini_adk_rs::tool::ToolPolicy::with_cache)
310 }
311
312 pub fn toolset(tools: Vec<Arc<dyn ToolFunction>>) -> ToolComposite {
314 ToolComposite {
315 entries: tools
316 .into_iter()
317 .map(ToolCompositeEntry::Function)
318 .collect(),
319 }
320 }
321
322 pub fn agent(
327 name: impl Into<String>,
328 description: impl Into<String>,
329 agent: impl TextAgent + 'static,
330 ) -> ToolComposite {
331 ToolComposite {
332 entries: vec![ToolCompositeEntry::Agent {
333 name: name.into(),
334 description: description.into(),
335 agent: Arc::new(agent),
336 }],
337 }
338 }
339
340 pub fn mcp(params: impl Into<String>) -> ToolComposite {
345 ToolComposite {
346 entries: vec![ToolCompositeEntry::Mcp {
347 params: params.into(),
348 }],
349 }
350 }
351
352 pub fn mock(
356 name: impl Into<String>,
357 description: impl Into<String>,
358 response: serde_json::Value,
359 ) -> ToolComposite {
360 ToolComposite {
361 entries: vec![ToolCompositeEntry::Mock {
362 name: name.into(),
363 description: description.into(),
364 response,
365 }],
366 }
367 }
368
369 pub fn schema(name: impl Into<String>, schema: serde_json::Value) -> ToolComposite {
373 ToolComposite {
374 entries: vec![ToolCompositeEntry::Schema {
375 name: name.into(),
376 schema,
377 }],
378 }
379 }
380
381 pub fn transform<F, Fut>(tool: ToolComposite, f: F) -> ToolComposite
386 where
387 F: Fn(serde_json::Value) -> Fut + Send + Sync + 'static,
388 Fut: Future<Output = serde_json::Value> + Send + 'static,
389 {
390 let f: TransformFn = Arc::new(
391 move |v: serde_json::Value| -> Pin<Box<dyn Future<Output = serde_json::Value> + Send>> {
392 Box::pin(f(v))
393 },
394 );
395 ToolComposite {
396 entries: tool
397 .entries
398 .into_iter()
399 .map(|entry| ToolCompositeEntry::Transform {
400 inner: Box::new(entry),
401 transformer: Arc::clone(&f),
402 })
403 .collect(),
404 }
405 }
406}
407
408#[derive(Clone, Debug)]
415pub enum DeferredTool {
416 Mcp {
418 params: String,
420 },
421}
422
423pub(crate) enum ToolResolution {
429 Runtime(Arc<dyn ToolFunction>),
431 BuiltIn(Tool),
433 Agent {
435 name: String,
437 description: String,
439 agent: Arc<dyn TextAgent>,
441 },
442 Deferred(DeferredTool),
444}
445
446impl ToolCompositeEntry {
447 #[cfg(test)]
448 fn classify_name(self) -> String {
449 match self.classify() {
450 ToolResolution::Runtime(f) => f.name().to_string(),
451 _ => String::new(),
452 }
453 }
454
455 pub(crate) fn classify(self) -> ToolResolution {
458 match self {
459 ToolCompositeEntry::Function(f) => ToolResolution::Runtime(f),
460 ToolCompositeEntry::BuiltIn(t) => ToolResolution::BuiltIn(t),
461 ToolCompositeEntry::Agent {
462 name,
463 description,
464 agent,
465 } => ToolResolution::Agent {
466 name,
467 description,
468 agent,
469 },
470 ToolCompositeEntry::Mock {
471 name,
472 description,
473 response,
474 } => ToolResolution::Runtime(Arc::new(SimpleTool::new(
475 name,
476 description,
477 None,
478 move |_args| {
479 let r = response.clone();
480 async move { Ok(r) }
481 },
482 ))),
483 ToolCompositeEntry::Transform { inner, transformer } => match inner.classify() {
484 ToolResolution::Runtime(f) => ToolResolution::Runtime(Arc::new(TransformTool {
485 inner: f,
486 transformer,
487 })),
488 other => other,
492 },
493 ToolCompositeEntry::Schema { name, schema } => {
494 ToolResolution::BuiltIn(Tool::functions(vec![FunctionDeclaration {
497 name,
498 description: String::new(),
499 parameters: Some(schema),
500 behavior: None,
501 }]))
502 }
503 ToolCompositeEntry::Mcp { params } => {
504 ToolResolution::Deferred(DeferredTool::Mcp { params })
505 }
506 }
507 }
508}
509
510struct TransformTool {
512 inner: Arc<dyn ToolFunction>,
513 #[allow(clippy::type_complexity)]
514 transformer: Arc<
515 dyn Fn(serde_json::Value) -> Pin<Box<dyn Future<Output = serde_json::Value> + Send>>
516 + Send
517 + Sync,
518 >,
519}
520
521#[async_trait::async_trait]
522impl ToolFunction for TransformTool {
523 fn name(&self) -> &str {
524 self.inner.name()
525 }
526
527 fn description(&self) -> &str {
528 self.inner.description()
529 }
530
531 fn parameters(&self) -> Option<serde_json::Value> {
532 self.inner.parameters()
533 }
534
535 async fn call(
536 &self,
537 args: serde_json::Value,
538 ) -> Result<serde_json::Value, gemini_adk_rs::error::ToolError> {
539 let result = self.inner.call(args).await?;
540 Ok((self.transformer)(result).await)
541 }
542
543 async fn call_with_context(
544 &self,
545 args: serde_json::Value,
546 ctx: gemini_adk_rs::tool::ToolContext,
547 ) -> Result<serde_json::Value, gemini_adk_rs::error::ToolError> {
548 let result = self.inner.call_with_context(args, ctx).await?;
549 Ok((self.transformer)(result).await)
550 }
551}
552
553#[cfg(test)]
554mod tests {
555 use super::*;
556
557 fn classify_one(c: ToolComposite) -> ToolResolution {
559 c.entries.into_iter().next().unwrap().classify()
560 }
561
562 #[test]
563 fn classify_maps_every_variant() {
564 assert!(matches!(
566 classify_one(T::mock("m", "d", serde_json::json!({"ok": true}))),
567 ToolResolution::Runtime(_)
568 ));
569 assert!(matches!(
570 classify_one(T::simple("s", "d", |a| async move { Ok(a) })),
571 ToolResolution::Runtime(_)
572 ));
573 assert!(matches!(
575 classify_one(T::google_search()),
576 ToolResolution::BuiltIn(_)
577 ));
578 assert!(matches!(
579 classify_one(T::schema("s", serde_json::json!({"type": "object"}))),
580 ToolResolution::BuiltIn(_)
581 ));
582 assert!(matches!(
584 classify_one(T::mcp("node ./server.js")),
585 ToolResolution::Deferred(DeferredTool::Mcp { .. })
586 ));
587 }
588
589 #[test]
590 fn a_single_tool_function_converts_into_a_composite() {
591 let composite: ToolComposite =
592 SimpleTool::new("one", "one", None, |_| async { Ok(serde_json::json!(1)) }).into();
593 assert_eq!(composite.len(), 1);
594 let arc: Arc<dyn ToolFunction> = Arc::new(SimpleTool::new("two", "two", None, |_| async {
595 Ok(serde_json::json!(2))
596 }));
597 let composite: ToolComposite = arc.into();
598 assert_eq!(composite.len(), 1);
599 assert_eq!(composite.entries[0].clone().classify_name(), "two");
600 }
601
602 #[tokio::test]
603 async fn mock_resolves_to_callable_runtime_tool() {
604 let resolution = classify_one(T::mock(
605 "weather",
606 "Mock weather",
607 serde_json::json!({"temp": 22}),
608 ));
609 let ToolResolution::Runtime(tool) = resolution else {
610 panic!("mock should resolve to a runtime tool");
611 };
612 assert_eq!(tool.name(), "weather");
613 let out = tool.call(serde_json::json!({})).await.unwrap();
614 assert_eq!(out, serde_json::json!({"temp": 22}));
615 }
616
617 #[tokio::test]
618 async fn transform_wraps_inner_runtime_result() {
619 let composite = T::transform(
620 T::mock("base", "d", serde_json::json!({"n": 1})),
621 |mut v| async move {
622 v["doubled"] = serde_json::json!(true);
623 v
624 },
625 );
626 let ToolResolution::Runtime(tool) = classify_one(composite) else {
627 panic!("transform over a mock should resolve to a runtime tool");
628 };
629 assert_eq!(tool.name(), "base");
630 let out = tool.call(serde_json::json!({})).await.unwrap();
631 assert_eq!(out, serde_json::json!({"n": 1, "doubled": true}));
632 }
633
634 #[test]
635 fn google_search_creates_composite() {
636 let t = T::google_search();
637 assert_eq!(t.len(), 1);
638 }
639
640 #[test]
641 fn url_context_creates_composite() {
642 let t = T::url_context();
643 assert_eq!(t.len(), 1);
644 }
645
646 #[test]
647 fn code_execution_creates_composite() {
648 let t = T::code_execution();
649 assert_eq!(t.len(), 1);
650 }
651
652 #[test]
653 fn compose_with_bitor() {
654 let t = T::google_search() + T::url_context() + T::code_execution();
655 assert_eq!(t.len(), 3);
656 }
657
658 #[test]
659 fn simple_creates_tool() {
660 let t = T::simple("greet", "Greets the user", |_args| async {
661 Ok(serde_json::json!({"message": "hello"}))
662 });
663 assert_eq!(t.len(), 1);
664 match &t.entries[0] {
665 ToolCompositeEntry::Function(f) => assert_eq!(f.name(), "greet"),
666 _ => panic!("expected Function entry"),
667 }
668 }
669
670 #[tokio::test]
671 async fn timeout_modifier_enforces_timeout() {
672 use gemini_adk_rs::ToolError;
673 use std::time::Duration;
674
675 let t = T::timeout(
676 T::simple("slow", "slow tool", |_| async move {
677 tokio::time::sleep(Duration::from_secs(3600)).await;
678 Ok(serde_json::json!({"ok": true}))
679 }),
680 Duration::from_millis(50),
681 );
682 match &t.entries[0] {
683 ToolCompositeEntry::Function(f) => match f.call(serde_json::json!({})).await {
684 Err(ToolError::Timeout(d)) => assert_eq!(d, Duration::from_millis(50)),
685 other => panic!("expected Timeout, got {other:?}"),
686 },
687 _ => panic!("expected Function entry"),
688 }
689 }
690
691 #[tokio::test]
692 async fn cached_modifier_memoizes_results() {
693 use std::sync::atomic::{AtomicU32, Ordering};
694
695 let counter = Arc::new(AtomicU32::new(0));
696 let c = counter.clone();
697 let t = T::cached(T::simple("count", "counts calls", move |_| {
698 let c = c.clone();
699 async move {
700 let n = c.fetch_add(1, Ordering::SeqCst) + 1;
701 Ok(serde_json::json!({"n": n}))
702 }
703 }));
704 match &t.entries[0] {
705 ToolCompositeEntry::Function(f) => {
706 let first = f.call(serde_json::json!({"x": 1})).await.unwrap();
707 let second = f.call(serde_json::json!({"x": 1})).await.unwrap();
708 assert_eq!(first, second);
709 assert_eq!(first["n"], 1);
710 assert_eq!(counter.load(Ordering::SeqCst), 1);
711 }
712 _ => panic!("expected Function entry"),
713 }
714 }
715
716 #[test]
717 fn confirm_modifier_wraps_function() {
718 let t = T::confirm(
721 T::simple("danger", "dangerous", |_| async move {
722 Ok(serde_json::json!({}))
723 }),
724 "are you sure?",
725 );
726 match &t.entries[0] {
727 ToolCompositeEntry::Function(f) => assert_eq!(f.name(), "danger"),
728 _ => panic!("expected Function entry"),
729 }
730 }
731
732 #[test]
733 fn toolset_combines_functions() {
734 let tool_a: Arc<dyn ToolFunction> =
735 Arc::new(SimpleTool::new("a", "tool a", None, |_| async {
736 Ok(serde_json::json!(null))
737 }));
738 let tool_b: Arc<dyn ToolFunction> =
739 Arc::new(SimpleTool::new("b", "tool b", None, |_| async {
740 Ok(serde_json::json!(null))
741 }));
742 let t = T::toolset(vec![tool_a, tool_b]);
743 assert_eq!(t.len(), 2);
744 }
745}