1use std::sync::Arc;
7
8use gemini_adk_rs::error::ConfigError;
9use gemini_adk_rs::llm::BaseLlm;
10use gemini_adk_rs::middleware::Middleware;
11use gemini_adk_rs::text::{LlmTextAgent, TextAgent};
12use gemini_adk_rs::tool::{ToolDispatcher, ToolFunction, ToolKind};
13use gemini_genai_rs::prelude::{Modality, ModelId, Tool, Voice};
14
15use crate::compose::context::ContextComposite;
16use crate::compose::guards::GuardComposite;
17use crate::compose::middleware::MiddlewareComposite;
18use crate::compose::tools::ToolComposite;
19
20type LlmProviderFn = Arc<dyn Fn(&gemini_adk_rs::State) -> Arc<dyn BaseLlm> + Send + Sync>;
21
22#[derive(Clone)]
24struct AgentBuilderInner {
25 name: String,
26 model: Option<ModelId>,
27 instruction: Option<String>,
28 instruction_provider: Option<Arc<dyn gemini_adk_rs::instruction::InstructionProvider>>,
29 llm_provider: Option<LlmProviderFn>,
30 voice: Option<Voice>,
31 temperature: Option<f32>,
32 top_p: Option<f32>,
33 top_k: Option<u32>,
34 max_output_tokens: Option<u32>,
35 stop_sequences: Vec<String>,
36 response_modalities: Option<Vec<Modality>>,
37 thinking_budget: Option<u32>,
38 tools: Vec<ToolEntry>,
39 built_in_tools: Vec<Tool>,
40 writes: Vec<String>,
41 reads: Vec<String>,
42 sub_agents: Vec<AgentBuilder>,
43 isolate: bool,
44 stay: bool,
45 description: Option<String>,
46 output_schema: Option<serde_json::Value>,
47 output_key: Option<String>,
48 transfer_to_agent: Option<String>,
49 middleware_layers: Vec<Arc<dyn Middleware>>,
51 confirmation_provider: Option<Arc<dyn gemini_adk_rs::confirmation::ConfirmationProvider>>,
53 config_errors: Vec<String>,
56}
57
58#[derive(Clone)]
60pub enum ToolEntry {
61 Runtime(Arc<dyn ToolEntryTrait>),
63 Declaration(Tool),
65}
66
67pub trait ToolEntryTrait: Send + Sync + 'static {
69 fn name(&self) -> &str;
71 fn to_tool_kind(&self) -> ToolKind;
73}
74
75#[derive(Clone)]
156pub struct AgentBuilder {
157 inner: Arc<AgentBuilderInner>,
158}
159
160impl AgentBuilder {
161 pub fn new(name: impl Into<String>) -> Self {
163 Self {
164 inner: Arc::new(AgentBuilderInner {
165 name: name.into(),
166 model: None,
167 instruction: None,
168 instruction_provider: None,
169 llm_provider: None,
170 voice: None,
171 temperature: None,
172 top_p: None,
173 top_k: None,
174 max_output_tokens: None,
175 stop_sequences: Vec::new(),
176 response_modalities: None,
177 thinking_budget: None,
178 tools: Vec::new(),
179 built_in_tools: Vec::new(),
180 writes: Vec::new(),
181 reads: Vec::new(),
182 sub_agents: Vec::new(),
183 isolate: false,
184 stay: false,
185 description: None,
186 output_schema: None,
187 output_key: None,
188 transfer_to_agent: None,
189 middleware_layers: Vec::new(),
190 confirmation_provider: None,
191 config_errors: Vec::new(),
192 }),
193 }
194 }
195
196 fn mutate(&self) -> AgentBuilderInner {
199 (*self.inner).clone()
200 }
201
202 fn with(inner: AgentBuilderInner) -> Self {
203 Self {
204 inner: Arc::new(inner),
205 }
206 }
207
208 pub fn name(&self) -> &str {
212 &self.inner.name
213 }
214
215 pub fn get_model(&self) -> Option<&ModelId> {
217 self.inner.model.as_ref()
218 }
219
220 pub fn get_instruction(&self) -> Option<&str> {
222 self.inner.instruction.as_deref()
223 }
224
225 pub fn get_voice(&self) -> Option<&Voice> {
227 self.inner.voice.as_ref()
228 }
229
230 pub fn get_temperature(&self) -> Option<f32> {
232 self.inner.temperature
233 }
234
235 pub fn is_text_only(&self) -> bool {
237 self.inner
238 .response_modalities
239 .as_ref()
240 .map(|m| m == &[Modality::Text])
241 .unwrap_or(false)
242 }
243
244 pub fn get_thinking_budget(&self) -> Option<u32> {
246 self.inner.thinking_budget
247 }
248
249 pub fn get_writes(&self) -> &[String] {
251 &self.inner.writes
252 }
253
254 pub fn get_reads(&self) -> &[String] {
256 &self.inner.reads
257 }
258
259 pub fn get_sub_agents(&self) -> &[AgentBuilder] {
261 &self.inner.sub_agents
262 }
263
264 pub fn is_isolated(&self) -> bool {
266 self.inner.isolate
267 }
268
269 pub fn is_stay(&self) -> bool {
271 self.inner.stay
272 }
273
274 pub fn tool_count(&self) -> usize {
276 self.inner.tools.len() + self.inner.built_in_tools.len()
277 }
278
279 pub fn get_top_p(&self) -> Option<f32> {
281 self.inner.top_p
282 }
283
284 pub fn get_top_k(&self) -> Option<u32> {
286 self.inner.top_k
287 }
288
289 pub fn get_max_output_tokens(&self) -> Option<u32> {
291 self.inner.max_output_tokens
292 }
293
294 pub fn get_stop_sequences(&self) -> &[String] {
296 &self.inner.stop_sequences
297 }
298
299 pub fn get_description(&self) -> Option<&str> {
301 self.inner.description.as_deref()
302 }
303
304 pub fn get_output_schema(&self) -> Option<&serde_json::Value> {
306 self.inner.output_schema.as_ref()
307 }
308
309 pub fn get_output_key(&self) -> Option<&str> {
311 self.inner.output_key.as_deref()
312 }
313
314 pub fn get_transfer_to(&self) -> Option<&str> {
316 self.inner.transfer_to_agent.as_deref()
317 }
318
319 pub fn middleware_layer_count(&self) -> usize {
321 self.inner.middleware_layers.len()
322 }
323
324 pub fn model(self, model: ModelId) -> Self {
328 let mut inner = self.mutate();
329 inner.model = Some(model);
330 Self::with(inner)
331 }
332
333 pub fn instruction(self, inst: impl Into<String>) -> Self {
335 let mut inner = self.mutate();
336 inner.instruction = Some(inst.into());
337 Self::with(inner)
338 }
339
340 pub fn instruction_provider(
345 self,
346 provider: impl gemini_adk_rs::instruction::InstructionProvider + 'static,
347 ) -> Self {
348 let mut inner = self.mutate();
349 inner.instruction_provider = Some(Arc::new(provider));
350 Self::with(inner)
351 }
352
353 pub fn llm_provider(
358 self,
359 provider: impl Fn(&gemini_adk_rs::State) -> Arc<dyn BaseLlm> + Send + Sync + 'static,
360 ) -> Self {
361 let mut inner = self.mutate();
362 inner.llm_provider = Some(Arc::new(provider));
363 Self::with(inner)
364 }
365
366 pub fn voice(self, voice: Voice) -> Self {
368 let mut inner = self.mutate();
369 inner.voice = Some(voice);
370 Self::with(inner)
371 }
372
373 pub fn temperature(self, t: f32) -> Self {
375 let mut inner = self.mutate();
376 inner.temperature = Some(t);
377 Self::with(inner)
378 }
379
380 pub fn text_only(self) -> Self {
382 let mut inner = self.mutate();
383 inner.response_modalities = Some(vec![Modality::Text]);
384 Self::with(inner)
385 }
386
387 pub fn response_modalities(self, modalities: Vec<Modality>) -> Self {
389 let mut inner = self.mutate();
390 inner.response_modalities = Some(modalities);
391 Self::with(inner)
392 }
393
394 pub fn thinking(self, budget: u32) -> Self {
396 let mut inner = self.mutate();
397 inner.thinking_budget = Some(budget);
398 Self::with(inner)
399 }
400
401 pub fn url_context(self) -> Self {
403 let mut inner = self.mutate();
404 inner.built_in_tools.push(Tool::url_context());
405 Self::with(inner)
406 }
407
408 pub fn google_search(self) -> Self {
410 let mut inner = self.mutate();
411 inner.built_in_tools.push(Tool::google_search());
412 Self::with(inner)
413 }
414
415 pub fn code_execution(self) -> Self {
417 let mut inner = self.mutate();
418 inner.built_in_tools.push(Tool::code_execution());
419 Self::with(inner)
420 }
421
422 pub fn writes(self, key: impl Into<String>) -> Self {
424 let mut inner = self.mutate();
425 inner.writes.push(key.into());
426 Self::with(inner)
427 }
428
429 pub fn reads(self, key: impl Into<String>) -> Self {
431 let mut inner = self.mutate();
432 inner.reads.push(key.into());
433 Self::with(inner)
434 }
435
436 pub fn artifacts(
444 self,
445 artifacts: impl Into<crate::compose::artifacts::ArtifactComposite>,
446 ) -> Self {
447 let artifacts = artifacts.into();
448 let mut builder = self;
449 for input in artifacts.all_inputs() {
450 builder = builder.reads(format!("artifact:{}", input.name));
451 }
452 for output in artifacts.all_outputs() {
453 builder = builder.writes(format!("artifact:{}", output.name));
454 }
455 builder
456 }
457
458 pub fn sub_agent(self, agent: AgentBuilder) -> Self {
460 let mut inner = self.mutate();
461 inner.sub_agents.push(agent);
462 Self::with(inner)
463 }
464
465 pub fn isolate(self) -> Self {
467 let mut inner = self.mutate();
468 inner.isolate = true;
469 Self::with(inner)
470 }
471
472 pub fn stay(self) -> Self {
474 let mut inner = self.mutate();
475 inner.stay = true;
476 Self::with(inner)
477 }
478
479 pub fn top_p(self, p: f32) -> Self {
481 let mut inner = self.mutate();
482 inner.top_p = Some(p);
483 Self::with(inner)
484 }
485
486 pub fn top_k(self, k: u32) -> Self {
488 let mut inner = self.mutate();
489 inner.top_k = Some(k);
490 Self::with(inner)
491 }
492
493 pub fn max_output_tokens(self, n: u32) -> Self {
495 let mut inner = self.mutate();
496 inner.max_output_tokens = Some(n);
497 Self::with(inner)
498 }
499
500 pub fn stop_sequences(self, seqs: Vec<String>) -> Self {
502 let mut inner = self.mutate();
503 inner.stop_sequences = seqs;
504 Self::with(inner)
505 }
506
507 pub fn description(self, desc: impl Into<String>) -> Self {
509 let mut inner = self.mutate();
510 inner.description = Some(desc.into());
511 Self::with(inner)
512 }
513
514 pub fn output_schema(self, schema: serde_json::Value) -> Self {
516 let mut inner = self.mutate();
517 inner.output_schema = Some(schema);
518 Self::with(inner)
519 }
520
521 pub fn output<T: schemars::JsonSchema>(self) -> Self {
528 self.output_schema(gemini_adk_rs::tool::wire_schema::<T>())
529 }
530
531 pub fn output_key(self, key: impl Into<String>) -> Self {
533 let mut inner = self.mutate();
534 inner.output_key = Some(key.into());
535 Self::with(inner)
536 }
537
538 pub fn transfer_to(self, agent_name: impl Into<String>) -> Self {
540 let mut inner = self.mutate();
541 inner.transfer_to_agent = Some(agent_name.into());
542 Self::with(inner)
543 }
544
545 #[deprecated(since = "2.1.0", note = "use `instruction`")]
549 pub fn instruct(self, inst: impl Into<String>) -> Self {
550 self.instruction(inst)
551 }
552
553 #[deprecated(since = "2.1.0", note = "use `description`")]
555 pub fn describe(self, desc: impl Into<String>) -> Self {
556 self.description(desc)
557 }
558
559 pub fn tool(self, f: impl ToolFunction + 'static) -> Self {
573 self.tools(ToolComposite::from_function(Arc::new(f)))
574 }
575
576 pub fn tools(self, tools: impl Into<ToolComposite>) -> Self {
591 use crate::compose::tools::{DeferredTool, ToolResolution};
592 let mut inner = self.mutate();
593 for entry in tools.into().entries {
594 match entry.classify() {
595 ToolResolution::Runtime(f) => {
596 inner
597 .tools
598 .push(ToolEntry::Runtime(Arc::new(ToolFunctionEntry(f))));
599 }
600 ToolResolution::BuiltIn(t) => {
601 inner.built_in_tools.push(t);
602 }
603 ToolResolution::Agent {
604 name,
605 description,
606 agent,
607 } => {
608 let tool = gemini_adk_rs::TextAgentTool::from_arc(
610 name,
611 description,
612 agent,
613 gemini_adk_rs::State::new(),
614 );
615 inner
616 .tools
617 .push(ToolEntry::Runtime(Arc::new(ToolFunctionEntry(Arc::new(
618 tool,
619 )))));
620 }
621 ToolResolution::Deferred(DeferredTool::Mcp { params }) => {
622 inner.config_errors.push(format!(
628 "T::mcp({params:?}) cannot be attached to a text AgentBuilder: MCP \
629 toolsets need an async connection, which only a Live session performs \
630 (`Live::builder().tools(T::mcp(..))`)"
631 ));
632 }
633 }
634 }
635 Self::with(inner)
636 }
637
638 pub fn guard(self, guard: impl Into<GuardComposite>) -> Self {
653 let mut inner = self.mutate();
654 inner.middleware_layers.push(guard.into().into_middleware());
655 Self::with(inner)
656 }
657
658 pub fn context(self, policy: impl Into<ContextComposite>) -> Self {
671 let mut inner = self.mutate();
672 inner
673 .middleware_layers
674 .push(policy.into().into_middleware());
675 Self::with(inner)
676 }
677
678 #[deprecated(since = "2.1.0", note = "use `isolate`, which this has always called")]
680 pub fn no_peers(self) -> Self {
681 self.isolate()
682 }
683
684 pub fn middleware(self, middleware: impl Into<MiddlewareComposite>) -> Self {
699 let mut inner = self.mutate();
700 inner.middleware_layers.extend(middleware.into().layers);
701 Self::with(inner)
702 }
703
704 pub fn confirmation_provider(
712 self,
713 provider: Arc<dyn gemini_adk_rs::confirmation::ConfirmationProvider>,
714 ) -> Self {
715 let mut inner = self.mutate();
716 inner.confirmation_provider = Some(provider);
717 Self::with(inner)
718 }
719
720 fn text_agent_issues(&self) -> Vec<String> {
728 let inner = &self.inner;
729 let mut issues = Vec::new();
730 if inner.voice.is_some() {
731 issues.push(
732 "`voice(..)` has no effect on a text agent, which returns text; \
733 remove it, or set the voice on `Live::builder()`"
734 .to_string(),
735 );
736 }
737 if inner
738 .response_modalities
739 .as_ref()
740 .is_some_and(|m| m.iter().any(|m| *m != Modality::Text))
741 {
742 issues.push(
743 "a text agent only produces text; audio and image `response_modalities` \
744 need a `Live::builder()` session"
745 .to_string(),
746 );
747 }
748 if let Some(model) = &inner.model
749 && gemini_adk_rs::llm::ModelCapabilities::infer_from_id(model.as_str()).live_bidi
750 {
751 issues.push(format!(
752 "`{model}` is a Live model, which the text API does not serve; use a text \
753 model such as `gemini-flash-latest`, or run it on `Live::builder()`"
754 ));
755 }
756 if inner.confirmation_provider.is_none() {
757 let gated: Vec<String> = inner
758 .tools
759 .iter()
760 .filter_map(|entry| match entry {
761 ToolEntry::Runtime(t) => match t.to_tool_kind() {
762 ToolKind::Function(f) if f.requires_confirmation() => {
763 Some(f.name().to_string())
764 }
765 _ => None,
766 },
767 ToolEntry::Declaration(_) => None,
768 })
769 .collect();
770 if !gated.is_empty() {
771 issues.push(format!(
772 "`{}` must be confirmed before running (`T::confirm`), but nothing can \
773 confirm it; set `.confirmation_provider(..)`",
774 gated.join("`, `")
775 ));
776 }
777 }
778 let transfer: Vec<&str> = [
779 (!inner.sub_agents.is_empty(), "sub_agent"),
780 (inner.transfer_to_agent.is_some(), "transfer_to"),
781 (inner.stay, "stay"),
782 (inner.isolate, "isolate/no_peers"),
783 ]
784 .into_iter()
785 .filter_map(|(set, name)| set.then_some(name))
786 .collect();
787 if !transfer.is_empty() {
788 issues.push(format!(
789 "`{}` configure agent transfer, which a text agent does not perform; \
790 compose text agents with `>>`, `|` and `T::agent(..)` instead",
791 transfer.join("`, `")
792 ));
793 }
794 issues
795 }
796
797 pub fn build(self, llm: impl BaseLlm + 'static) -> Result<Arc<dyn TextAgent>, ConfigError> {
829 let mut issues = self.inner.config_errors.clone();
830 issues.extend(self.text_agent_issues());
831 if !issues.is_empty() {
832 return Err(ConfigError { issues });
833 }
834 let inner = &self.inner;
835 let mut agent = LlmTextAgent::new(&inner.name, llm);
836
837 if let Some(inst) = &inner.instruction {
838 agent = agent.instruction(inst);
839 }
840 if let Some(provider) = &inner.instruction_provider {
841 agent = agent.instruction_provider(provider.clone());
842 }
843 if let Some(provider) = &inner.llm_provider {
844 let provider_clone = provider.clone();
845 agent = agent.llm_provider(move |state| provider_clone(state));
846 }
847 if let Some(model) = &inner.model {
848 agent = agent.model(model.as_str());
849 }
850 if let Some(t) = inner.temperature {
851 agent = agent.temperature(t);
852 }
853 if let Some(p) = inner.top_p {
854 agent = agent.top_p(p);
855 }
856 if let Some(k) = inner.top_k {
857 agent = agent.top_k(k);
858 }
859 if let Some(n) = inner.max_output_tokens {
860 agent = agent.max_output_tokens(n);
861 }
862 if !inner.stop_sequences.is_empty() {
863 agent = agent.stop_sequences(inner.stop_sequences.iter().cloned());
864 }
865 if let Some(budget) = inner.thinking_budget {
866 agent = agent.thinking_budget(budget);
867 }
868 if let Some(schema) = &inner.output_schema {
869 agent = agent.response_schema(schema.clone());
870 }
871 if let Some(key) = &inner.output_key {
872 agent = agent.output_key(key);
873 }
874 let declarations = inner.tools.iter().filter_map(|entry| match entry {
875 ToolEntry::Declaration(tool) => Some(tool),
876 ToolEntry::Runtime(_) => None,
877 });
878 for tool in inner.built_in_tools.iter().chain(declarations) {
879 agent = agent.built_in_tool(tool.clone());
880 }
881
882 if !self.inner.tools.is_empty() {
884 let mut dispatcher = ToolDispatcher::new();
885 for entry in &self.inner.tools {
886 match entry {
887 ToolEntry::Runtime(t) => {
888 let kind = t.to_tool_kind();
889 match kind {
890 ToolKind::Function(f) => dispatcher.register_function(f),
891 ToolKind::Streaming(s) => dispatcher.register_streaming(s),
892 ToolKind::InputStream(i) => dispatcher.register_input_streaming(i),
893 }
894 }
895 ToolEntry::Declaration(_) => {}
897 }
898 }
899 if let Some(provider) = &self.inner.confirmation_provider {
900 dispatcher.set_confirmation_provider(provider.clone());
901 }
902 if !dispatcher.is_empty() {
903 agent = agent.tools(Arc::new(dispatcher));
904 }
905 }
906
907 for mw in &self.inner.middleware_layers {
909 agent = agent.add_middleware(mw.clone());
910 }
911
912 Ok(Arc::new(agent))
913 }
914}
915
916#[derive(Clone)]
918struct ToolFunctionEntry(Arc<dyn ToolFunction>);
919
920impl ToolEntryTrait for ToolFunctionEntry {
921 fn name(&self) -> &str {
922 self.0.name()
923 }
924
925 fn to_tool_kind(&self) -> ToolKind {
926 ToolKind::Function(self.0.clone())
927 }
928}
929
930impl std::fmt::Debug for AgentBuilder {
931 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
932 f.debug_struct("AgentBuilder")
933 .field("name", &self.inner.name)
934 .field("model", &self.inner.model)
935 .field("instruction", &self.inner.instruction)
936 .field("temperature", &self.inner.temperature)
937 .field("text_only", &self.is_text_only())
938 .field("tool_count", &self.tool_count())
939 .field("sub_agents", &self.inner.sub_agents.len())
940 .finish()
941 }
942}
943
944#[cfg(test)]
945mod tests {
946 use super::*;
947 use async_trait::async_trait;
948 use gemini_adk_rs::llm::{LlmError, LlmRequest, LlmResponse};
949 use gemini_genai_rs::prelude::{Content, Part, Role};
950
951 struct MockLlm(String);
953
954 #[async_trait]
955 impl BaseLlm for MockLlm {
956 fn model_id(&self) -> &str {
957 "mock"
958 }
959 async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
960 Ok(LlmResponse {
961 content: Content {
962 role: Some(Role::Model),
963 parts: vec![Part::Text {
964 text: self.0.clone(),
965 }],
966 },
967 finish_reason: Some("STOP".into()),
968 usage: None,
969 })
970 }
971 }
972
973 #[test]
974 fn builder_creates_with_name() {
975 let b = AgentBuilder::new("test-agent");
976 assert_eq!(b.name(), "test-agent");
977 }
978
979 #[test]
980 fn fluent_chaining_works() {
981 let b = AgentBuilder::new("agent")
982 .instruction("Be helpful")
983 .temperature(0.7)
984 .model(ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO);
985
986 assert_eq!(b.get_instruction(), Some("Be helpful"));
987 assert_eq!(b.get_temperature(), Some(0.7));
988 assert_eq!(b.get_model(), Some(&ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO));
989 }
990
991 #[test]
992 fn copy_on_write_clone_independence() {
993 let base = AgentBuilder::new("base").temperature(0.5);
994 let variant = base.clone().temperature(0.9);
995
996 assert_eq!(base.get_temperature(), Some(0.5));
998 assert_eq!(variant.get_temperature(), Some(0.9));
1000 }
1001
1002 #[test]
1003 fn text_only_sets_modalities() {
1004 let b = AgentBuilder::new("text").text_only();
1005 assert!(b.is_text_only());
1006 }
1007
1008 #[test]
1009 fn url_context_adds_tool() {
1010 let b = AgentBuilder::new("search").url_context();
1011 assert_eq!(b.tool_count(), 1);
1012 }
1013
1014 #[test]
1015 fn google_search_adds_tool() {
1016 let b = AgentBuilder::new("search").google_search();
1017 assert_eq!(b.tool_count(), 1);
1018 }
1019
1020 #[test]
1021 fn code_execution_adds_tool() {
1022 let b = AgentBuilder::new("code").code_execution();
1023 assert_eq!(b.tool_count(), 1);
1024 }
1025
1026 #[test]
1027 fn thinking_sets_budget() {
1028 let b = AgentBuilder::new("thinker").thinking(2048);
1029 assert_eq!(b.get_thinking_budget(), Some(2048));
1030 }
1031
1032 #[test]
1033 fn writes_and_reads_keys() {
1034 let b = AgentBuilder::new("data").writes("output").reads("input");
1035 assert_eq!(b.get_writes(), &["output"]);
1036 assert_eq!(b.get_reads(), &["input"]);
1037 }
1038
1039 #[test]
1040 fn sub_agent_registration() {
1041 let child = AgentBuilder::new("child");
1042 let parent = AgentBuilder::new("parent").sub_agent(child);
1043 assert_eq!(parent.get_sub_agents().len(), 1);
1044 assert_eq!(parent.get_sub_agents()[0].name(), "child");
1045 }
1046
1047 #[test]
1048 fn isolate_and_stay() {
1049 let b = AgentBuilder::new("agent").isolate().stay();
1050 assert!(b.is_isolated());
1051 assert!(b.is_stay());
1052 }
1053
1054 #[test]
1055 fn debug_display() {
1056 let b = AgentBuilder::new("debug-test");
1057 let debug = format!("{b:?}");
1058 assert!(debug.contains("debug-test"));
1059 }
1060
1061 #[test]
1062 fn top_p_sets_value() {
1063 let b = AgentBuilder::new("agent").top_p(0.95);
1064 assert_eq!(b.get_top_p(), Some(0.95));
1065 }
1066
1067 #[test]
1068 fn top_k_sets_value() {
1069 let b = AgentBuilder::new("agent").top_k(40);
1070 assert_eq!(b.get_top_k(), Some(40));
1071 }
1072
1073 #[test]
1074 fn max_output_tokens_sets_value() {
1075 let b = AgentBuilder::new("agent").max_output_tokens(4096);
1076 assert_eq!(b.get_max_output_tokens(), Some(4096));
1077 }
1078
1079 #[test]
1080 fn stop_sequences_sets_value() {
1081 let b =
1082 AgentBuilder::new("agent").stop_sequences(vec!["END".to_string(), "STOP".to_string()]);
1083 assert_eq!(b.get_stop_sequences().len(), 2);
1084 }
1085
1086 #[test]
1087 fn description_sets_value() {
1088 let b = AgentBuilder::new("agent").description("A helpful agent");
1089 assert_eq!(b.get_description(), Some("A helpful agent"));
1090 }
1091
1092 #[test]
1093 fn output_schema_sets_value() {
1094 let schema = serde_json::json!({"type": "object"});
1095 let b = AgentBuilder::new("agent").output_schema(schema.clone());
1096 assert_eq!(b.get_output_schema(), Some(&schema));
1097 }
1098
1099 #[test]
1100 fn transfer_to_sets_value() {
1101 let b = AgentBuilder::new("agent").transfer_to("target-agent");
1102 assert_eq!(b.get_transfer_to(), Some("target-agent"));
1103 }
1104
1105 #[test]
1106 fn full_fluent_chain() {
1107 let b = AgentBuilder::new("full-agent")
1108 .model(ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO)
1109 .instruction("Be helpful")
1110 .temperature(0.7)
1111 .top_p(0.95)
1112 .top_k(40)
1113 .max_output_tokens(4096)
1114 .thinking(2048)
1115 .description("A fully configured agent")
1116 .google_search()
1117 .writes("output")
1118 .reads("input");
1119
1120 assert_eq!(b.name(), "full-agent");
1121 assert_eq!(b.get_temperature(), Some(0.7));
1122 assert_eq!(b.get_top_p(), Some(0.95));
1123 assert_eq!(b.get_top_k(), Some(40));
1124 assert_eq!(b.get_max_output_tokens(), Some(4096));
1125 assert_eq!(b.get_thinking_budget(), Some(2048));
1126 assert_eq!(b.get_description(), Some("A fully configured agent"));
1127 assert_eq!(b.tool_count(), 1);
1128 }
1129
1130 #[tokio::test]
1133 async fn build_produces_executable_agent() {
1134 let llm: Arc<dyn BaseLlm> = Arc::new(MockLlm("built agent output".into()));
1135 let agent = AgentBuilder::new("test")
1136 .instruction("Be helpful")
1137 .temperature(0.5)
1138 .build(llm)
1139 .unwrap();
1140
1141 assert_eq!(agent.name(), "test");
1142 let state = gemini_adk_rs::State::new();
1143 let result = agent.run(&state).await.unwrap();
1144 assert_eq!(result, "built agent output");
1145 }
1146
1147 #[tokio::test]
1148 async fn build_stores_output_in_state() {
1149 let llm: Arc<dyn BaseLlm> = Arc::new(MockLlm("state output".into()));
1150 let agent = AgentBuilder::new("test").build(llm).unwrap();
1151 let state = gemini_adk_rs::State::new();
1152 agent.run(&state).await.unwrap();
1153 assert_eq!(state.get::<String>("output"), Some("state output".into()));
1154 }
1155
1156 #[tokio::test]
1157 async fn build_reads_input_from_state() {
1158 use gemini_adk_rs::llm::LlmRequest;
1159
1160 struct EchoLlm;
1162 #[async_trait]
1163 impl BaseLlm for EchoLlm {
1164 fn model_id(&self) -> &str {
1165 "echo"
1166 }
1167 async fn generate(&self, req: LlmRequest) -> Result<LlmResponse, LlmError> {
1168 let text: String = req
1169 .contents
1170 .iter()
1171 .flat_map(|c| &c.parts)
1172 .filter_map(|p| match p {
1173 Part::Text { text } => Some(text.as_str()),
1174 _ => None,
1175 })
1176 .collect::<Vec<_>>()
1177 .join("");
1178 Ok(LlmResponse {
1179 content: Content {
1180 role: Some(Role::Model),
1181 parts: vec![Part::Text { text }],
1182 },
1183 finish_reason: Some("STOP".into()),
1184 usage: None,
1185 })
1186 }
1187 }
1188
1189 let agent = AgentBuilder::new("echo").build(Arc::new(EchoLlm)).unwrap();
1190 let state = gemini_adk_rs::State::new();
1191 let _ = state.set("input", "hello from state");
1192 let result = agent.run(&state).await.unwrap();
1193 assert!(result.contains("hello from state"));
1194 }
1195
1196 struct ToolCallingMockLlm {
1200 tool_name: &'static str,
1201 final_text: &'static str,
1202 }
1203
1204 #[async_trait]
1205 impl BaseLlm for ToolCallingMockLlm {
1206 fn model_id(&self) -> &str {
1207 "tool-mock"
1208 }
1209
1210 async fn generate(&self, req: LlmRequest) -> Result<LlmResponse, LlmError> {
1211 use gemini_genai_rs::prelude::FunctionCall;
1212
1213 let already_responded = req
1215 .contents
1216 .iter()
1217 .flat_map(|c| &c.parts)
1218 .any(|p| matches!(p, Part::FunctionResponse { .. }));
1219
1220 if already_responded {
1221 Ok(LlmResponse {
1222 content: Content {
1223 role: Some(Role::Model),
1224 parts: vec![Part::Text {
1225 text: self.final_text.to_string(),
1226 }],
1227 },
1228 finish_reason: Some("STOP".into()),
1229 usage: None,
1230 })
1231 } else {
1232 Ok(LlmResponse {
1233 content: Content {
1234 role: Some(Role::Model),
1235 parts: vec![Part::FunctionCall {
1236 function_call: FunctionCall {
1237 name: self.tool_name.to_string(),
1238 args: serde_json::json!({"x": 1}),
1239 id: Some("call-1".into()),
1240 },
1241 }],
1242 },
1243 finish_reason: None,
1244 usage: None,
1245 })
1246 }
1247 }
1248 }
1249
1250 #[tokio::test]
1252 async fn middleware_hooks_fire_end_to_end() {
1253 use crate::compose::middleware::M;
1254 use gemini_adk_rs::tool::SimpleTool;
1255 use std::sync::atomic::{AtomicUsize, Ordering};
1256
1257 let before_model_count = Arc::new(AtomicUsize::new(0));
1258 let after_tool_count = Arc::new(AtomicUsize::new(0));
1259
1260 let bm = before_model_count.clone();
1261 let at = after_tool_count.clone();
1262
1263 let mw = M::before_model(move |_req| {
1264 bm.fetch_add(1, Ordering::SeqCst);
1265 Ok(())
1266 }) >> M::after_tool(move |_call, _result| {
1267 at.fetch_add(1, Ordering::SeqCst);
1268 Ok(())
1269 });
1270
1271 let llm: Arc<dyn BaseLlm> = Arc::new(ToolCallingMockLlm {
1272 tool_name: "echo_tool",
1273 final_text: "done",
1274 });
1275
1276 let agent = AgentBuilder::new("mw-test")
1277 .middleware(mw)
1278 .tool(SimpleTool::new(
1279 "echo_tool",
1280 "Echo tool",
1281 None,
1282 |_args| async move { Ok(serde_json::json!({"echo": true})) },
1283 ))
1284 .build(llm)
1285 .unwrap();
1286
1287 let state = gemini_adk_rs::State::new();
1288 let result = agent.run(&state).await.unwrap();
1289 assert_eq!(result, "done");
1290
1291 assert_eq!(
1293 before_model_count.load(Ordering::SeqCst),
1294 2,
1295 "before_model should fire for each generate() call"
1296 );
1297 assert_eq!(
1299 after_tool_count.load(Ordering::SeqCst),
1300 1,
1301 "after_tool should fire once for the tool dispatch"
1302 );
1303 }
1304
1305 #[test]
1307 fn middleware_copy_on_write() {
1308 use crate::compose::middleware::M;
1309
1310 let base = AgentBuilder::new("base").instruction("base");
1311 let with_mw = base.clone().middleware(M::log() >> M::latency());
1312
1313 assert_eq!(base.middleware_layer_count(), 0);
1315 assert_eq!(with_mw.middleware_layer_count(), 2);
1317 }
1318
1319 #[tokio::test]
1321 async fn middleware_on_error_fires_on_failure() {
1322 use crate::compose::middleware::M;
1323 use gemini_adk_rs::llm::LlmError;
1324 use std::sync::atomic::{AtomicUsize, Ordering};
1325
1326 let error_count = Arc::new(AtomicUsize::new(0));
1327 let ec = error_count.clone();
1328
1329 let mw = M::on_error(move |_err| {
1330 ec.fetch_add(1, Ordering::SeqCst);
1331 Ok(())
1332 });
1333
1334 struct FailLlm;
1335 #[async_trait]
1336 impl BaseLlm for FailLlm {
1337 fn model_id(&self) -> &str {
1338 "fail"
1339 }
1340 async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
1341 Err(LlmError::RequestFailed("boom".into()))
1342 }
1343 }
1344
1345 let agent = AgentBuilder::new("error-test")
1346 .middleware(mw)
1347 .build(Arc::new(FailLlm))
1348 .unwrap();
1349
1350 let state = gemini_adk_rs::State::new();
1351 let result = agent.run(&state).await;
1352 assert!(result.is_err(), "agent should fail");
1353 assert_eq!(
1354 error_count.load(Ordering::SeqCst),
1355 1,
1356 "on_error should fire exactly once"
1357 );
1358 }
1359
1360 struct RecordingLlm {
1365 text: &'static str,
1366 seen_len: Arc<std::sync::atomic::AtomicUsize>,
1367 }
1368
1369 #[async_trait]
1370 impl BaseLlm for RecordingLlm {
1371 fn model_id(&self) -> &str {
1372 "recording-mock"
1373 }
1374
1375 async fn generate(&self, req: LlmRequest) -> Result<LlmResponse, LlmError> {
1376 self.seen_len
1377 .store(req.contents.len(), std::sync::atomic::Ordering::SeqCst);
1378 Ok(LlmResponse {
1379 content: Content {
1380 role: Some(Role::Model),
1381 parts: vec![Part::Text {
1382 text: self.text.to_string(),
1383 }],
1384 },
1385 finish_reason: Some("STOP".into()),
1386 usage: None,
1387 })
1388 }
1389 }
1390
1391 #[tokio::test]
1392 async fn guard_blocks_violating_output() {
1393 use crate::compose::guards::G;
1394
1395 let llm: Arc<dyn BaseLlm> = Arc::new(RecordingLlm {
1396 text: "you can reach me at agent@example.com",
1397 seen_len: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
1398 });
1399
1400 let agent = AgentBuilder::new("guarded")
1401 .guard(G::pii())
1402 .build(llm)
1403 .unwrap();
1404
1405 let state = gemini_adk_rs::State::new();
1406 let err = agent.run(&state).await.unwrap_err();
1407 assert!(
1408 err.to_string().contains("guard violation"),
1409 "PII guard should veto the response, got: {err}"
1410 );
1411 }
1412
1413 #[tokio::test]
1414 async fn guard_allows_clean_output() {
1415 use crate::compose::guards::G;
1416
1417 let llm: Arc<dyn BaseLlm> = Arc::new(RecordingLlm {
1418 text: "all clean here",
1419 seen_len: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
1420 });
1421
1422 let agent = AgentBuilder::new("guarded")
1423 .guard(G::pii() + G::length(1, 1000))
1424 .build(llm)
1425 .unwrap();
1426
1427 let state = gemini_adk_rs::State::new();
1428 let result = agent.run(&state).await.unwrap();
1429 assert_eq!(result, "all clean here");
1430 }
1431
1432 #[tokio::test]
1433 async fn context_policy_rewrites_request_history() {
1434 use crate::compose::context::C;
1435
1436 let seen = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1439 let llm: Arc<dyn BaseLlm> = Arc::new(RecordingLlm {
1440 text: "ok",
1441 seen_len: seen.clone(),
1442 });
1443
1444 let agent = AgentBuilder::new("ctx")
1445 .context(C::prepend(Content::user("system preamble")))
1446 .build(llm)
1447 .unwrap();
1448
1449 let state = gemini_adk_rs::State::new();
1450 let _ = state.set("input", "hello");
1451 let _ = agent.run(&state).await.unwrap();
1452 assert_eq!(
1453 seen.load(std::sync::atomic::Ordering::SeqCst),
1454 2,
1455 "context policy should have prepended a turn before the model call"
1456 );
1457 }
1458
1459 #[tokio::test]
1460 async fn context_window_trims_history() {
1461 use crate::compose::context::C;
1462
1463 let seen = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1466 let llm: Arc<dyn BaseLlm> = Arc::new(RecordingLlm {
1467 text: "ok",
1468 seen_len: seen.clone(),
1469 });
1470
1471 let agent = AgentBuilder::new("ctx")
1472 .context(
1473 C::prepend(Content::user("a")) >> C::prepend(Content::user("b")) >> C::window(1),
1474 )
1475 .build(llm)
1476 .unwrap();
1477
1478 let state = gemini_adk_rs::State::new();
1479 let _ = state.set("input", "hello");
1480 let _ = agent.run(&state).await.unwrap();
1481 assert_eq!(
1482 seen.load(std::sync::atomic::Ordering::SeqCst),
1483 1,
1484 "window(1) should trim history to the last turn"
1485 );
1486 }
1487}
1488
1489#[cfg(test)]
1490mod mcp_rejection_tests {
1491 use super::*;
1492 use gemini_adk_rs::llm::{LlmError, LlmRequest, LlmResponse};
1493
1494 struct NeverLlm;
1495 #[async_trait::async_trait]
1496 impl BaseLlm for NeverLlm {
1497 fn model_id(&self) -> &str {
1498 "never"
1499 }
1500 async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
1501 Err(LlmError::RequestFailed("never".into()))
1502 }
1503 }
1504
1505 #[test]
1508 fn mcp_toolset_is_a_build_error() {
1509 use crate::compose::tools::T;
1510 let err = AgentBuilder::new("text")
1511 .tools(T::mcp("node ./server.js"))
1512 .build(Arc::new(NeverLlm))
1513 .err()
1514 .expect("build must fail");
1515 assert!(err.to_string().contains("T::mcp"), "{err}");
1516 assert!(err.to_string().contains("Live"), "{err}");
1517 }
1518
1519 #[test]
1521 fn tool_accepts_values_and_arcs() {
1522 use gemini_adk_rs::tool::SimpleTool;
1523 let arc: Arc<dyn ToolFunction> = Arc::new(SimpleTool::new("a", "a", None, |_| async {
1524 Ok(serde_json::json!({}))
1525 }));
1526 let b = AgentBuilder::new("t")
1527 .tool(SimpleTool::new("b", "b", None, |_| async {
1528 Ok(serde_json::json!({}))
1529 }))
1530 .tool(arc.clone())
1531 .tools(arc);
1532 assert_eq!(b.tool_count(), 3);
1533 }
1534}
1535
1536#[cfg(test)]
1537mod honest_build_tests {
1538 use super::*;
1539 use gemini_adk_rs::llm::{LlmResponse, MockLlm};
1540
1541 #[tokio::test]
1545 async fn every_setting_reaches_the_request() {
1546 let llm = MockLlm::script([LlmResponse::from_text(r#"{"city":"Paris"}"#)]);
1547 let agent = AgentBuilder::new("configured")
1548 .model(ModelId::new("gemini-2.5-pro"))
1549 .instruction("Answer in JSON.")
1550 .temperature(0.1)
1551 .top_p(0.8)
1552 .top_k(10)
1553 .max_output_tokens(128)
1554 .stop_sequences(vec!["END".into()])
1555 .thinking(256)
1556 .output_schema(serde_json::json!({ "type": "object" }))
1557 .output_key("answer")
1558 .google_search()
1559 .build(llm.clone())
1560 .expect("a text agent can honour all of these");
1561
1562 let state = gemini_adk_rs::State::new();
1563 state.set("input", "Capital of France?").unwrap();
1564 agent.run(&state).await.unwrap();
1565
1566 let sent = llm.last_request().unwrap();
1567 assert_eq!(sent.model.as_deref(), Some("gemini-2.5-pro"));
1568 assert_eq!(sent.system_instruction.as_deref(), Some("Answer in JSON."));
1569 assert_eq!(sent.temperature, Some(0.1));
1570 assert_eq!(sent.top_p, Some(0.8));
1571 assert_eq!(sent.top_k, Some(10));
1572 assert_eq!(sent.max_output_tokens, Some(128));
1573 assert_eq!(sent.stop_sequences, ["END"]);
1574 assert_eq!(sent.thinking_budget, Some(256));
1575 assert_eq!(sent.response_mime_type.as_deref(), Some("application/json"));
1576 assert_eq!(
1577 sent.response_json_schema,
1578 Some(serde_json::json!({ "type": "object" }))
1579 );
1580 assert_eq!(sent.tools, vec![Tool::google_search()]);
1581 assert_eq!(
1582 state.get::<String>("answer").as_deref(),
1583 Some(r#"{"city":"Paris"}"#)
1584 );
1585 }
1586
1587 #[tokio::test]
1589 async fn built_in_and_function_tools_are_both_declared() {
1590 use crate::compose::tools::T;
1591 let llm = MockLlm::text("ok");
1592 let agent = AgentBuilder::new("both")
1593 .tools(
1594 T::simple("ping", "Ping", |_| async { Ok(serde_json::json!({})) })
1595 + T::code_execution(),
1596 )
1597 .url_context()
1598 .build(llm.clone())
1599 .unwrap();
1600 agent.run(&gemini_adk_rs::State::new()).await.unwrap();
1601 let tools = llm.last_request().unwrap().tools;
1602 assert!(tools.contains(&Tool::url_context()), "{tools:?}");
1603 assert!(tools.contains(&Tool::code_execution()), "{tools:?}");
1604 assert!(
1605 tools.iter().any(|t| t
1606 .function_declarations
1607 .as_ref()
1608 .is_some_and(|d| d.iter().any(|f| f.name == "ping"))),
1609 "{tools:?}"
1610 );
1611 }
1612
1613 #[test]
1616 fn live_only_settings_are_build_errors() {
1617 let err = AgentBuilder::new("confused")
1618 .model(ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO)
1619 .voice(Voice::Kore)
1620 .response_modalities(vec![Modality::Audio])
1621 .sub_agent(AgentBuilder::new("child"))
1622 .transfer_to("child")
1623 .build(MockLlm::text("unused"))
1624 .err()
1625 .expect("build must fail");
1626 assert_eq!(err.issues.len(), 4, "{err}");
1627 let message = err.to_string();
1628 for needle in [
1629 "voice",
1630 "response_modalities",
1631 "Live model",
1632 "sub_agent",
1633 "transfer_to",
1634 ] {
1635 assert!(message.contains(needle), "missing {needle}: {message}");
1636 }
1637 }
1638
1639 fn refund_tool() -> crate::compose::tools::ToolComposite {
1640 use crate::compose::tools::T;
1641 T::confirm(
1642 T::simple("refund", "Refund the order", |_| async {
1643 Ok(serde_json::json!({ "refunded": true }))
1644 }),
1645 "Refunds move money",
1646 )
1647 }
1648
1649 #[test]
1652 fn a_confirm_tool_without_a_provider_is_a_build_error() {
1653 let err = AgentBuilder::new("support")
1654 .tools(refund_tool())
1655 .build(MockLlm::text("x"))
1656 .err()
1657 .expect("build must fail");
1658 assert!(err.to_string().contains("refund"), "{err}");
1659 assert!(err.to_string().contains("confirmation_provider"), "{err}");
1660 }
1661
1662 #[tokio::test]
1664 async fn a_declined_call_tells_the_model_why() {
1665 use gemini_adk_rs::confirmation::StaticConfirmation;
1666 use gemini_genai_rs::prelude::Part;
1667
1668 let llm = MockLlm::script([
1669 LlmResponse::tool_call("refund", serde_json::json!({})),
1670 LlmResponse::from_text("I could not refund that."),
1671 ]);
1672 let agent = AgentBuilder::new("support")
1673 .tools(refund_tool())
1674 .confirmation_provider(StaticConfirmation::deny_all("over the refund limit"))
1675 .build(llm.clone())
1676 .unwrap();
1677 agent.run(&gemini_adk_rs::State::new()).await.unwrap();
1678
1679 let returned = llm
1680 .last_request()
1681 .unwrap()
1682 .contents
1683 .iter()
1684 .flat_map(|c| c.parts.clone())
1685 .find_map(|p| match p {
1686 Part::FunctionResponse { function_response } => Some(function_response.response),
1687 _ => None,
1688 })
1689 .expect("the refusal is sent back");
1690 let error = returned["error"].as_str().unwrap_or_default();
1691 assert!(error.contains("over the refund limit"), "{returned}");
1692 assert!(returned.get("refunded").is_none(), "the tool must not run");
1693 }
1694
1695 #[tokio::test]
1697 async fn output_type_constrains_the_reply() {
1698 #[derive(serde::Deserialize, schemars::JsonSchema)]
1699 struct City {
1700 name: String,
1701 }
1702 let llm = MockLlm::text(r#"{"name":"Paris"}"#);
1703 let agent = AgentBuilder::new("geo")
1704 .output::<City>()
1705 .build(llm.clone())
1706 .unwrap();
1707 let result = agent
1708 .run_with(
1709 gemini_adk_rs::text::RunRequest::new("Capital of France?"),
1710 &gemini_adk_rs::State::new(),
1711 )
1712 .await
1713 .unwrap();
1714 assert_eq!(result.parse::<City>().unwrap().name, "Paris");
1715 let schema = llm.last_request().unwrap().response_json_schema.unwrap();
1716 assert_eq!(schema["properties"]["name"]["type"], "string");
1717 }
1718
1719 #[test]
1720 fn text_only_is_a_text_agent_setting() {
1721 assert!(
1722 AgentBuilder::new("t")
1723 .text_only()
1724 .build(MockLlm::text("x"))
1725 .is_ok()
1726 );
1727 }
1728}