gemini_adk_fluent_rs/
builder.rs

1//! AgentBuilder — copy-on-write immutable builder for fluent agent construction.
2//!
3//! Every mutation returns a new builder (original unchanged), so builders
4//! are safely shareable as templates.
5
6use std::sync::Arc;
7
8use gemini_adk_rs::llm::BaseLlm;
9use gemini_adk_rs::text::{LlmTextAgent, TextAgent};
10use gemini_adk_rs::tool::{ToolDispatcher, ToolFunction, ToolKind};
11use gemini_genai_rs::prelude::{GeminiModel, Modality, Tool, Voice};
12
13use crate::compose::context::ContextPolicy;
14use crate::compose::guards::GComposite;
15use crate::compose::tools::ToolComposite;
16
17/// Inner state of an AgentBuilder — shared via Arc for copy-on-write.
18#[derive(Clone)]
19struct AgentBuilderInner {
20    name: String,
21    model: Option<GeminiModel>,
22    instruction: Option<String>,
23    voice: Option<Voice>,
24    temperature: Option<f32>,
25    top_p: Option<f32>,
26    top_k: Option<u32>,
27    max_output_tokens: Option<u32>,
28    stop_sequences: Vec<String>,
29    response_modalities: Option<Vec<Modality>>,
30    thinking_budget: Option<u32>,
31    tools: Vec<ToolEntry>,
32    built_in_tools: Vec<Tool>,
33    writes: Vec<String>,
34    reads: Vec<String>,
35    sub_agents: Vec<AgentBuilder>,
36    isolate: bool,
37    stay: bool,
38    description: Option<String>,
39    output_schema: Option<serde_json::Value>,
40    output_key: Option<String>,
41    transfer_to_agent: Option<String>,
42}
43
44/// An entry in the builder's tool list — either a runtime ToolKind or a declaration.
45#[derive(Clone)]
46pub enum ToolEntry {
47    /// A runtime tool with a handler function.
48    Runtime(Arc<dyn ToolEntryTrait>),
49    /// A wire-level tool declaration (e.g., built-in tools like Google Search).
50    Declaration(Tool),
51}
52
53/// Trait for tool entries that can provide a name (for dedup/inspection).
54pub trait ToolEntryTrait: Send + Sync + 'static {
55    /// The tool's registered name.
56    fn name(&self) -> &str;
57    /// Convert this entry into the runtime `ToolKind` variant for dispatch.
58    fn to_tool_kind(&self) -> ToolKind;
59}
60
61/// Alias for [`AgentBuilder`] — matches upstream Python `Agent("name")` naming.
62pub type Agent = AgentBuilder;
63
64/// Copy-on-write immutable builder for agent construction.
65///
66/// Every setter returns a new `AgentBuilder`, leaving the original unchanged.
67/// This makes builders safe to share as templates.
68///
69/// # Basic Usage
70///
71/// ```rust
72/// use gemini_adk_fluent_rs::builder::AgentBuilder;
73/// use gemini_genai_rs::prelude::GeminiModel;
74///
75/// let agent = AgentBuilder::new("analyst")
76///     .model(GeminiModel::Gemini2_0FlashLive)
77///     .instruction("Analyze the given topic")
78///     .temperature(0.3);
79///
80/// assert_eq!(agent.name(), "analyst");
81/// assert_eq!(agent.get_temperature(), Some(0.3));
82/// ```
83///
84/// # Copy-on-Write Pattern
85///
86/// Cloning a builder and modifying the clone leaves the original unchanged.
87/// This is useful for creating template builders with shared defaults.
88///
89/// ```rust
90/// use gemini_adk_fluent_rs::builder::AgentBuilder;
91///
92/// let base = AgentBuilder::new("researcher")
93///     .instruction("You are a research assistant.")
94///     .temperature(0.5);
95///
96/// let creative = base.clone().temperature(0.9);
97/// let precise  = base.clone().temperature(0.1);
98///
99/// // Original unchanged
100/// assert_eq!(base.get_temperature(), Some(0.5));
101/// assert_eq!(creative.get_temperature(), Some(0.9));
102/// assert_eq!(precise.get_temperature(), Some(0.1));
103/// ```
104///
105/// # Sampling Parameters
106///
107/// ```rust
108/// use gemini_adk_fluent_rs::builder::AgentBuilder;
109///
110/// let agent = AgentBuilder::new("sampler")
111///     .temperature(0.7)
112///     .top_p(0.95)
113///     .top_k(40)
114///     .max_output_tokens(4096);
115///
116/// assert_eq!(agent.get_top_p(), Some(0.95));
117/// assert_eq!(agent.get_top_k(), Some(40));
118/// assert_eq!(agent.get_max_output_tokens(), Some(4096));
119/// ```
120///
121/// # Built-in Tools
122///
123/// ```rust
124/// use gemini_adk_fluent_rs::builder::AgentBuilder;
125///
126/// let agent = AgentBuilder::new("searcher")
127///     .google_search()
128///     .code_execution()
129///     .url_context();
130///
131/// assert_eq!(agent.tool_count(), 3);
132/// ```
133///
134/// # Thinking Budget
135///
136/// ```rust
137/// use gemini_adk_fluent_rs::builder::AgentBuilder;
138///
139/// let agent = AgentBuilder::new("thinker")
140///     .thinking(2048);
141///
142/// assert_eq!(agent.get_thinking_budget(), Some(2048));
143/// ```
144#[derive(Clone)]
145pub struct AgentBuilder {
146    inner: Arc<AgentBuilderInner>,
147}
148
149impl AgentBuilder {
150    /// Create a new builder with the given agent name.
151    pub fn new(name: impl Into<String>) -> Self {
152        Self {
153            inner: Arc::new(AgentBuilderInner {
154                name: name.into(),
155                model: None,
156                instruction: None,
157                voice: None,
158                temperature: None,
159                top_p: None,
160                top_k: None,
161                max_output_tokens: None,
162                stop_sequences: Vec::new(),
163                response_modalities: None,
164                thinking_budget: None,
165                tools: Vec::new(),
166                built_in_tools: Vec::new(),
167                writes: Vec::new(),
168                reads: Vec::new(),
169                sub_agents: Vec::new(),
170                isolate: false,
171                stay: false,
172                description: None,
173                output_schema: None,
174                output_key: None,
175                transfer_to_agent: None,
176            }),
177        }
178    }
179
180    // ── Private helper: clone-on-write ──
181
182    fn mutate(&self) -> AgentBuilderInner {
183        (*self.inner).clone()
184    }
185
186    fn with(inner: AgentBuilderInner) -> Self {
187        Self {
188            inner: Arc::new(inner),
189        }
190    }
191
192    // ── Accessors ──
193
194    /// The agent name.
195    pub fn name(&self) -> &str {
196        &self.inner.name
197    }
198
199    /// Configured model, if any.
200    pub fn get_model(&self) -> Option<&GeminiModel> {
201        self.inner.model.as_ref()
202    }
203
204    /// Configured instruction, if any.
205    pub fn get_instruction(&self) -> Option<&str> {
206        self.inner.instruction.as_deref()
207    }
208
209    /// Configured voice, if any.
210    pub fn get_voice(&self) -> Option<&Voice> {
211        self.inner.voice.as_ref()
212    }
213
214    /// Configured temperature, if any.
215    pub fn get_temperature(&self) -> Option<f32> {
216        self.inner.temperature
217    }
218
219    /// Whether text-only mode is set.
220    pub fn is_text_only(&self) -> bool {
221        self.inner
222            .response_modalities
223            .as_ref()
224            .map(|m| m == &[Modality::Text])
225            .unwrap_or(false)
226    }
227
228    /// Configured thinking budget, if any.
229    pub fn get_thinking_budget(&self) -> Option<u32> {
230        self.inner.thinking_budget
231    }
232
233    /// State keys this agent writes.
234    pub fn get_writes(&self) -> &[String] {
235        &self.inner.writes
236    }
237
238    /// State keys this agent reads.
239    pub fn get_reads(&self) -> &[String] {
240        &self.inner.reads
241    }
242
243    /// Sub-agents registered.
244    pub fn get_sub_agents(&self) -> &[AgentBuilder] {
245        &self.inner.sub_agents
246    }
247
248    /// Whether agent runs in isolated state.
249    pub fn is_isolated(&self) -> bool {
250        self.inner.isolate
251    }
252
253    /// Whether agent stays after transfer.
254    pub fn is_stay(&self) -> bool {
255        self.inner.stay
256    }
257
258    /// Number of tool entries.
259    pub fn tool_count(&self) -> usize {
260        self.inner.tools.len() + self.inner.built_in_tools.len()
261    }
262
263    /// Configured top_p, if any.
264    pub fn get_top_p(&self) -> Option<f32> {
265        self.inner.top_p
266    }
267
268    /// Configured top_k, if any.
269    pub fn get_top_k(&self) -> Option<u32> {
270        self.inner.top_k
271    }
272
273    /// Configured max_output_tokens, if any.
274    pub fn get_max_output_tokens(&self) -> Option<u32> {
275        self.inner.max_output_tokens
276    }
277
278    /// Configured stop sequences.
279    pub fn get_stop_sequences(&self) -> &[String] {
280        &self.inner.stop_sequences
281    }
282
283    /// Configured description, if any.
284    pub fn get_description(&self) -> Option<&str> {
285        self.inner.description.as_deref()
286    }
287
288    /// Configured output schema, if any.
289    pub fn get_output_schema(&self) -> Option<&serde_json::Value> {
290        self.inner.output_schema.as_ref()
291    }
292
293    /// Get the configured output key.
294    pub fn get_output_key(&self) -> Option<&str> {
295        self.inner.output_key.as_deref()
296    }
297
298    /// Configured transfer target agent, if any.
299    pub fn get_transfer_to(&self) -> Option<&str> {
300        self.inner.transfer_to_agent.as_deref()
301    }
302
303    // ── Fluent Setters (copy-on-write) ──
304
305    /// Set the Gemini model.
306    pub fn model(self, model: GeminiModel) -> Self {
307        let mut inner = self.mutate();
308        inner.model = Some(model);
309        Self::with(inner)
310    }
311
312    /// Set the system instruction.
313    pub fn instruction(self, inst: impl Into<String>) -> Self {
314        let mut inner = self.mutate();
315        inner.instruction = Some(inst.into());
316        Self::with(inner)
317    }
318
319    /// Set the output voice.
320    pub fn voice(self, voice: Voice) -> Self {
321        let mut inner = self.mutate();
322        inner.voice = Some(voice);
323        Self::with(inner)
324    }
325
326    /// Set the temperature.
327    pub fn temperature(self, t: f32) -> Self {
328        let mut inner = self.mutate();
329        inner.temperature = Some(t);
330        Self::with(inner)
331    }
332
333    /// Set text-only mode (no audio output).
334    pub fn text_only(self) -> Self {
335        let mut inner = self.mutate();
336        inner.response_modalities = Some(vec![Modality::Text]);
337        Self::with(inner)
338    }
339
340    /// Set response modalities explicitly.
341    pub fn response_modalities(self, modalities: Vec<Modality>) -> Self {
342        let mut inner = self.mutate();
343        inner.response_modalities = Some(modalities);
344        Self::with(inner)
345    }
346
347    /// Enable thinking with a token budget.
348    pub fn thinking(self, budget: u32) -> Self {
349        let mut inner = self.mutate();
350        inner.thinking_budget = Some(budget);
351        Self::with(inner)
352    }
353
354    /// Add a built-in URL context tool.
355    pub fn url_context(self) -> Self {
356        let mut inner = self.mutate();
357        inner.built_in_tools.push(Tool::url_context());
358        Self::with(inner)
359    }
360
361    /// Add a built-in Google Search tool.
362    pub fn google_search(self) -> Self {
363        let mut inner = self.mutate();
364        inner.built_in_tools.push(Tool::google_search());
365        Self::with(inner)
366    }
367
368    /// Add a built-in code execution tool.
369    pub fn code_execution(self) -> Self {
370        let mut inner = self.mutate();
371        inner.built_in_tools.push(Tool::code_execution());
372        Self::with(inner)
373    }
374
375    /// Declare a state key this agent writes.
376    pub fn writes(self, key: impl Into<String>) -> Self {
377        let mut inner = self.mutate();
378        inner.writes.push(key.into());
379        Self::with(inner)
380    }
381
382    /// Declare a state key this agent reads.
383    pub fn reads(self, key: impl Into<String>) -> Self {
384        let mut inner = self.mutate();
385        inner.reads.push(key.into());
386        Self::with(inner)
387    }
388
389    /// Add a sub-agent for transfer.
390    pub fn sub_agent(self, agent: AgentBuilder) -> Self {
391        let mut inner = self.mutate();
392        inner.sub_agents.push(agent);
393        Self::with(inner)
394    }
395
396    /// Run this agent in isolated state (no shared state).
397    pub fn isolate(self) -> Self {
398        let mut inner = self.mutate();
399        inner.isolate = true;
400        Self::with(inner)
401    }
402
403    /// Keep this agent active after transfer (don't tear down).
404    pub fn stay(self) -> Self {
405        let mut inner = self.mutate();
406        inner.stay = true;
407        Self::with(inner)
408    }
409
410    /// Set top_p (nucleus sampling).
411    pub fn top_p(self, p: f32) -> Self {
412        let mut inner = self.mutate();
413        inner.top_p = Some(p);
414        Self::with(inner)
415    }
416
417    /// Set top_k (top-k sampling).
418    pub fn top_k(self, k: u32) -> Self {
419        let mut inner = self.mutate();
420        inner.top_k = Some(k);
421        Self::with(inner)
422    }
423
424    /// Set maximum output tokens.
425    pub fn max_output_tokens(self, n: u32) -> Self {
426        let mut inner = self.mutate();
427        inner.max_output_tokens = Some(n);
428        Self::with(inner)
429    }
430
431    /// Set stop sequences.
432    pub fn stop_sequences(self, seqs: Vec<String>) -> Self {
433        let mut inner = self.mutate();
434        inner.stop_sequences = seqs;
435        Self::with(inner)
436    }
437
438    /// Set a description for this agent (used in tool/agent metadata).
439    pub fn description(self, desc: impl Into<String>) -> Self {
440        let mut inner = self.mutate();
441        inner.description = Some(desc.into());
442        Self::with(inner)
443    }
444
445    /// Set a JSON schema for structured output.
446    pub fn output_schema(self, schema: serde_json::Value) -> Self {
447        let mut inner = self.mutate();
448        inner.output_schema = Some(schema);
449        Self::with(inner)
450    }
451
452    /// Set the output key — agent's final text response is auto-saved to this state key.
453    pub fn output_key(self, key: impl Into<String>) -> Self {
454        let mut inner = self.mutate();
455        inner.output_key = Some(key.into());
456        Self::with(inner)
457    }
458
459    /// Set a default transfer target agent.
460    pub fn transfer_to(self, agent_name: impl Into<String>) -> Self {
461        let mut inner = self.mutate();
462        inner.transfer_to_agent = Some(agent_name.into());
463        Self::with(inner)
464    }
465
466    // ── Upstream naming aliases ──
467
468    /// Alias for [`instruction`](Self::instruction) — matches upstream Python `Agent.instruct()`.
469    pub fn instruct(self, inst: impl Into<String>) -> Self {
470        self.instruction(inst)
471    }
472
473    /// Alias for [`description`](Self::description) — matches upstream Python `Agent.describe()`.
474    pub fn describe(self, desc: impl Into<String>) -> Self {
475        self.description(desc)
476    }
477
478    /// Register a single tool function.
479    ///
480    /// ```ignore
481    /// Agent::new("assistant").tool(Arc::new(my_tool))
482    /// ```
483    pub fn tool(self, f: Arc<dyn ToolFunction>) -> Self {
484        let mut inner = self.mutate();
485        inner
486            .tools
487            .push(ToolEntry::Runtime(Arc::new(ToolFunctionEntry(f))));
488        Self::with(inner)
489    }
490
491    /// Register multiple tools from a [`ToolComposite`].
492    ///
493    /// ```ignore
494    /// let tools = T::simple("greet", "Greet", |_| async { Ok(json!({})) })
495    ///     | T::google_search();
496    /// Agent::new("assistant").tools(tools)
497    /// ```
498    pub fn tools(self, composite: ToolComposite) -> Self {
499        use crate::compose::tools::ToolCompositeEntry;
500        let mut inner = self.mutate();
501        for entry in composite.entries {
502            match entry {
503                ToolCompositeEntry::Function(f) => {
504                    inner
505                        .tools
506                        .push(ToolEntry::Runtime(Arc::new(ToolFunctionEntry(f))));
507                }
508                ToolCompositeEntry::BuiltIn(t) => {
509                    inner.built_in_tools.push(t);
510                }
511                // Placeholder variants — not yet wired into the text agent builder.
512                _ => {
513                    // Agent, Mcp, A2a, Mock, OpenApi, Search, Schema, Transform
514                    // are currently only handled at the Live layer.
515                }
516            }
517        }
518        Self::with(inner)
519    }
520
521    /// Set a guard composite for output validation.
522    ///
523    /// Guards are evaluated after each agent response. This stores the guard
524    /// configuration for use at compile time.
525    pub fn guard(self, _guard: GComposite) -> Self {
526        // Guards are declarative metadata — stored for compile-time wiring.
527        // Full enforcement requires runtime integration (Phase 4+).
528        self
529    }
530
531    /// Set a context policy for conversation history management.
532    pub fn context(self, _policy: ContextPolicy) -> Self {
533        // Context policies are declarative — stored for compile-time wiring.
534        self
535    }
536
537    /// Disallow transfer to peer agents.
538    pub fn no_peers(self) -> Self {
539        self.isolate()
540    }
541
542    // ── Compilation ──
543
544    /// Compile this builder into an executable `TextAgent`.
545    ///
546    /// The LLM is required because `TextAgent` makes `BaseLlm::generate()` calls.
547    /// Builder configuration (instruction, temperature, tools) is transferred to
548    /// the resulting agent.
549    ///
550    /// ```rust,ignore
551    /// let agent = AgentBuilder::new("analyst")
552    ///     .instruction("Analyze the topic")
553    ///     .temperature(0.3)
554    ///     .build(llm);
555    ///
556    /// let result = agent.run(&state).await?;
557    /// ```
558    pub fn build(self, llm: Arc<dyn BaseLlm>) -> Arc<dyn TextAgent> {
559        let mut agent = LlmTextAgent::new(&self.inner.name, llm);
560
561        if let Some(inst) = &self.inner.instruction {
562            agent = agent.instruction(inst);
563        }
564        if let Some(t) = self.inner.temperature {
565            agent = agent.temperature(t);
566        }
567        if let Some(n) = self.inner.max_output_tokens {
568            agent = agent.max_output_tokens(n);
569        }
570
571        // Build ToolDispatcher from registered tools.
572        if !self.inner.tools.is_empty() {
573            let mut dispatcher = ToolDispatcher::new();
574            for entry in &self.inner.tools {
575                match entry {
576                    ToolEntry::Runtime(t) => {
577                        let kind = t.to_tool_kind();
578                        match kind {
579                            ToolKind::Function(f) => dispatcher.register_function(f),
580                            ToolKind::Streaming(s) => dispatcher.register_streaming(s),
581                            ToolKind::InputStream(i) => dispatcher.register_input_streaming(i),
582                        }
583                    }
584                    ToolEntry::Declaration(_) => {
585                        // Built-in tool declarations (google_search, etc.) are sent
586                        // as-is; they don't have runtime handlers for text dispatch.
587                    }
588                }
589            }
590            if !dispatcher.is_empty() {
591                agent = agent.tools(Arc::new(dispatcher));
592            }
593        }
594
595        Arc::new(agent)
596    }
597}
598
599/// Adapter that wraps an `Arc<dyn ToolFunction>` as a `ToolEntryTrait`.
600#[derive(Clone)]
601struct ToolFunctionEntry(Arc<dyn ToolFunction>);
602
603impl ToolEntryTrait for ToolFunctionEntry {
604    fn name(&self) -> &str {
605        self.0.name()
606    }
607
608    fn to_tool_kind(&self) -> ToolKind {
609        ToolKind::Function(self.0.clone())
610    }
611}
612
613impl std::fmt::Debug for AgentBuilder {
614    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
615        f.debug_struct("AgentBuilder")
616            .field("name", &self.inner.name)
617            .field("model", &self.inner.model)
618            .field("instruction", &self.inner.instruction)
619            .field("temperature", &self.inner.temperature)
620            .field("text_only", &self.is_text_only())
621            .field("tool_count", &self.tool_count())
622            .field("sub_agents", &self.inner.sub_agents.len())
623            .finish()
624    }
625}
626
627#[cfg(test)]
628mod tests {
629    use super::*;
630    use async_trait::async_trait;
631    use gemini_adk_rs::llm::{LlmError, LlmRequest, LlmResponse};
632    use gemini_genai_rs::prelude::{Content, Part, Role};
633
634    /// A mock LLM for build() tests.
635    struct MockLlm(String);
636
637    #[async_trait]
638    impl BaseLlm for MockLlm {
639        fn model_id(&self) -> &str {
640            "mock"
641        }
642        async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
643            Ok(LlmResponse {
644                content: Content {
645                    role: Some(Role::Model),
646                    parts: vec![Part::Text {
647                        text: self.0.clone(),
648                    }],
649                },
650                finish_reason: Some("STOP".into()),
651                usage: None,
652            })
653        }
654    }
655
656    #[test]
657    fn builder_creates_with_name() {
658        let b = AgentBuilder::new("test-agent");
659        assert_eq!(b.name(), "test-agent");
660    }
661
662    #[test]
663    fn fluent_chaining_works() {
664        let b = AgentBuilder::new("agent")
665            .instruction("Be helpful")
666            .temperature(0.7)
667            .model(GeminiModel::Gemini2_0FlashLive);
668
669        assert_eq!(b.get_instruction(), Some("Be helpful"));
670        assert_eq!(b.get_temperature(), Some(0.7));
671        assert_eq!(b.get_model(), Some(&GeminiModel::Gemini2_0FlashLive));
672    }
673
674    #[test]
675    fn copy_on_write_clone_independence() {
676        let base = AgentBuilder::new("base").temperature(0.5);
677        let variant = base.clone().temperature(0.9);
678
679        // Original unchanged
680        assert_eq!(base.get_temperature(), Some(0.5));
681        // Variant has new value
682        assert_eq!(variant.get_temperature(), Some(0.9));
683    }
684
685    #[test]
686    fn text_only_sets_modalities() {
687        let b = AgentBuilder::new("text").text_only();
688        assert!(b.is_text_only());
689    }
690
691    #[test]
692    fn url_context_adds_tool() {
693        let b = AgentBuilder::new("search").url_context();
694        assert_eq!(b.tool_count(), 1);
695    }
696
697    #[test]
698    fn google_search_adds_tool() {
699        let b = AgentBuilder::new("search").google_search();
700        assert_eq!(b.tool_count(), 1);
701    }
702
703    #[test]
704    fn code_execution_adds_tool() {
705        let b = AgentBuilder::new("code").code_execution();
706        assert_eq!(b.tool_count(), 1);
707    }
708
709    #[test]
710    fn thinking_sets_budget() {
711        let b = AgentBuilder::new("thinker").thinking(2048);
712        assert_eq!(b.get_thinking_budget(), Some(2048));
713    }
714
715    #[test]
716    fn writes_and_reads_keys() {
717        let b = AgentBuilder::new("data").writes("output").reads("input");
718        assert_eq!(b.get_writes(), &["output"]);
719        assert_eq!(b.get_reads(), &["input"]);
720    }
721
722    #[test]
723    fn sub_agent_registration() {
724        let child = AgentBuilder::new("child");
725        let parent = AgentBuilder::new("parent").sub_agent(child);
726        assert_eq!(parent.get_sub_agents().len(), 1);
727        assert_eq!(parent.get_sub_agents()[0].name(), "child");
728    }
729
730    #[test]
731    fn isolate_and_stay() {
732        let b = AgentBuilder::new("agent").isolate().stay();
733        assert!(b.is_isolated());
734        assert!(b.is_stay());
735    }
736
737    #[test]
738    fn debug_display() {
739        let b = AgentBuilder::new("debug-test");
740        let debug = format!("{:?}", b);
741        assert!(debug.contains("debug-test"));
742    }
743
744    #[test]
745    fn top_p_sets_value() {
746        let b = AgentBuilder::new("agent").top_p(0.95);
747        assert_eq!(b.get_top_p(), Some(0.95));
748    }
749
750    #[test]
751    fn top_k_sets_value() {
752        let b = AgentBuilder::new("agent").top_k(40);
753        assert_eq!(b.get_top_k(), Some(40));
754    }
755
756    #[test]
757    fn max_output_tokens_sets_value() {
758        let b = AgentBuilder::new("agent").max_output_tokens(4096);
759        assert_eq!(b.get_max_output_tokens(), Some(4096));
760    }
761
762    #[test]
763    fn stop_sequences_sets_value() {
764        let b =
765            AgentBuilder::new("agent").stop_sequences(vec!["END".to_string(), "STOP".to_string()]);
766        assert_eq!(b.get_stop_sequences().len(), 2);
767    }
768
769    #[test]
770    fn description_sets_value() {
771        let b = AgentBuilder::new("agent").description("A helpful agent");
772        assert_eq!(b.get_description(), Some("A helpful agent"));
773    }
774
775    #[test]
776    fn output_schema_sets_value() {
777        let schema = serde_json::json!({"type": "object"});
778        let b = AgentBuilder::new("agent").output_schema(schema.clone());
779        assert_eq!(b.get_output_schema(), Some(&schema));
780    }
781
782    #[test]
783    fn transfer_to_sets_value() {
784        let b = AgentBuilder::new("agent").transfer_to("target-agent");
785        assert_eq!(b.get_transfer_to(), Some("target-agent"));
786    }
787
788    #[test]
789    fn full_fluent_chain() {
790        let b = AgentBuilder::new("full-agent")
791            .model(GeminiModel::Gemini2_0FlashLive)
792            .instruction("Be helpful")
793            .temperature(0.7)
794            .top_p(0.95)
795            .top_k(40)
796            .max_output_tokens(4096)
797            .thinking(2048)
798            .description("A fully configured agent")
799            .google_search()
800            .writes("output")
801            .reads("input");
802
803        assert_eq!(b.name(), "full-agent");
804        assert_eq!(b.get_temperature(), Some(0.7));
805        assert_eq!(b.get_top_p(), Some(0.95));
806        assert_eq!(b.get_top_k(), Some(40));
807        assert_eq!(b.get_max_output_tokens(), Some(4096));
808        assert_eq!(b.get_thinking_budget(), Some(2048));
809        assert_eq!(b.get_description(), Some("A fully configured agent"));
810        assert_eq!(b.tool_count(), 1);
811    }
812
813    // ── build() tests ──
814
815    #[tokio::test]
816    async fn build_produces_executable_agent() {
817        let llm: Arc<dyn BaseLlm> = Arc::new(MockLlm("built agent output".into()));
818        let agent = AgentBuilder::new("test")
819            .instruction("Be helpful")
820            .temperature(0.5)
821            .build(llm);
822
823        assert_eq!(agent.name(), "test");
824        let state = gemini_adk_rs::State::new();
825        let result = agent.run(&state).await.unwrap();
826        assert_eq!(result, "built agent output");
827    }
828
829    #[tokio::test]
830    async fn build_stores_output_in_state() {
831        let llm: Arc<dyn BaseLlm> = Arc::new(MockLlm("state output".into()));
832        let agent = AgentBuilder::new("test").build(llm);
833        let state = gemini_adk_rs::State::new();
834        agent.run(&state).await.unwrap();
835        assert_eq!(state.get::<String>("output"), Some("state output".into()));
836    }
837
838    #[tokio::test]
839    async fn build_reads_input_from_state() {
840        use gemini_adk_rs::llm::LlmRequest;
841
842        // An LLM that echoes whatever it receives.
843        struct EchoLlm;
844        #[async_trait]
845        impl BaseLlm for EchoLlm {
846            fn model_id(&self) -> &str {
847                "echo"
848            }
849            async fn generate(&self, req: LlmRequest) -> Result<LlmResponse, LlmError> {
850                let text: String = req
851                    .contents
852                    .iter()
853                    .flat_map(|c| &c.parts)
854                    .filter_map(|p| match p {
855                        Part::Text { text } => Some(text.as_str()),
856                        _ => None,
857                    })
858                    .collect::<Vec<_>>()
859                    .join("");
860                Ok(LlmResponse {
861                    content: Content {
862                        role: Some(Role::Model),
863                        parts: vec![Part::Text { text }],
864                    },
865                    finish_reason: Some("STOP".into()),
866                    usage: None,
867                })
868            }
869        }
870
871        let agent = AgentBuilder::new("echo").build(Arc::new(EchoLlm));
872        let state = gemini_adk_rs::State::new();
873        state.set("input", "hello from state");
874        let result = agent.run(&state).await.unwrap();
875        assert!(result.contains("hello from state"));
876    }
877}
gemini_adk_fluent_rs/builder.rs

gemini_adk_fluent_rs/
builder.rs
