gemini_adk_fluent_rs/
operators.rs

1//! Operator algebra for agent composition.
2//!
3//! All types implementing `Composable` participate in the algebra:
4//!
5//! | Operator | Meaning            | Example                    |
6//! |----------|--------------------|----------------------------|
7//! | `>>`     | Sequential pipeline| `agent_a >> agent_b`       |
8//! | `\|`     | Parallel fan-out   | `agent_a \| agent_b`       |
9//! | `*`      | Loop (fixed)       | `agent * 3`                |
10//! | `//`     | Fallback chain     | `agent_a // agent_b`       |
11
12use std::sync::Arc;
13
14use gemini_adk_rs::error::ConfigError;
15use gemini_adk_rs::llm::BaseLlm;
16use gemini_adk_rs::middleware::{Middleware, MiddlewareChain};
17use gemini_adk_rs::text::{
18    FallbackTextAgent, LoopTextAgent, ParallelTextAgent, SequentialTextAgent, TextAgent,
19};
20
21use crate::builder::AgentBuilder;
22use crate::compose::middleware::MiddlewareComposite;
23
24/// A composable workflow node — can be sequenced, fan-out, looped, etc.
25#[derive(Clone, Debug)]
26pub enum Composable {
27    /// A single agent node.
28    Agent(AgentBuilder),
29    /// A sequential pipeline of steps.
30    Pipeline(Pipeline),
31    /// A parallel fan-out of branches.
32    FanOut(FanOut),
33    /// A loop with optional termination predicate.
34    Loop(Loop),
35    /// A fallback chain (try each until one succeeds).
36    Fallback(Fallback),
37    /// One agent applied to every item of a state list
38    /// ([`patterns::map_over`](crate::patterns::map_over)).
39    MapOver(crate::patterns::MapOver),
40}
41
42/// Sequential pipeline: execute steps in order, passing state between them.
43#[derive(Clone, Debug, Default)]
44pub struct Pipeline {
45    /// Ordered steps to execute sequentially.
46    pub steps: Vec<Composable>,
47    /// Name given to the compiled agent (default `"pipeline"`).
48    pub name: Option<String>,
49    /// Human-readable description, shown in `Debug` output.
50    pub description: Option<String>,
51}
52
53/// Parallel fan-out: execute branches concurrently, merge results.
54#[derive(Clone, Debug, Default)]
55pub struct FanOut {
56    /// Branches to execute concurrently.
57    pub branches: Vec<Composable>,
58    /// Name given to the compiled agent (default `"fan_out"`).
59    pub name: Option<String>,
60    /// Human-readable description, shown in `Debug` output.
61    pub description: Option<String>,
62}
63
64/// Loop: repeat an agent or pipeline up to `max` times, or until a predicate.
65#[derive(Clone)]
66pub struct Loop {
67    /// The composable to repeat.
68    pub body: Box<Composable>,
69    /// Maximum number of iterations.
70    pub max: u32,
71    /// Optional early-exit predicate evaluated after each iteration.
72    pub until: Option<LoopPredicate>,
73    /// Middleware attached to the loop agent (e.g. `M::on_loop` observers).
74    /// Set via [`Loop::middleware`] / [`Composable::middleware`]; construct as
75    /// `Vec::new()` in literals.
76    #[doc(hidden)]
77    pub middleware: Vec<Arc<dyn Middleware>>,
78    /// Name given to the compiled agent (default `"loop"`).
79    pub name: Option<String>,
80    /// Human-readable description, shown in `Debug` output.
81    pub description: Option<String>,
82}
83
84/// Predicate for conditional loop termination.
85#[derive(Clone)]
86pub struct LoopPredicate {
87    predicate: std::sync::Arc<dyn Fn(&serde_json::Value) -> bool + Send + Sync>,
88}
89
90impl LoopPredicate {
91    /// Create a new predicate from a closure that checks loop state.
92    pub fn new(f: impl Fn(&serde_json::Value) -> bool + Send + Sync + 'static) -> Self {
93        Self {
94            predicate: std::sync::Arc::new(f),
95        }
96    }
97
98    /// Evaluate the predicate against the current state.
99    pub fn check(&self, state: &serde_json::Value) -> bool {
100        (self.predicate)(state)
101    }
102}
103
104impl std::fmt::Debug for LoopPredicate {
105    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
106        f.write_str("LoopPredicate(<fn>)")
107    }
108}
109
110impl std::fmt::Debug for Loop {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        f.debug_struct("Loop")
113            .field("body", &self.body)
114            .field("max", &self.max)
115            .field("until", &self.until)
116            .field("name", &self.name)
117            .field("description", &self.description)
118            .finish()
119    }
120}
121
122/// Fallback chain: try each agent in sequence until one succeeds.
123#[derive(Clone)]
124pub struct Fallback {
125    /// Candidate composables tried in order until one succeeds.
126    pub candidates: Vec<Composable>,
127    /// Middleware attached to the fallback agent (e.g. `M::on_fallback`).
128    middleware: Vec<Arc<dyn Middleware>>,
129}
130
131impl std::fmt::Debug for Fallback {
132    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
133        f.debug_struct("Fallback")
134            .field("candidates", &self.candidates)
135            .finish()
136    }
137}
138
139/// Create a conditional loop predicate.
140pub fn until(
141    predicate: impl Fn(&serde_json::Value) -> bool + Send + Sync + 'static,
142) -> LoopPredicate {
143    LoopPredicate::new(predicate)
144}
145
146// ── Conversions ──
147
148impl From<AgentBuilder> for Composable {
149    fn from(b: AgentBuilder) -> Self {
150        Composable::Agent(b)
151    }
152}
153
154impl From<Pipeline> for Composable {
155    fn from(p: Pipeline) -> Self {
156        Composable::Pipeline(p)
157    }
158}
159
160impl From<FanOut> for Composable {
161    fn from(f: FanOut) -> Self {
162        Composable::FanOut(f)
163    }
164}
165
166impl From<Loop> for Composable {
167    fn from(l: Loop) -> Self {
168        Composable::Loop(l)
169    }
170}
171
172impl From<Fallback> for Composable {
173    fn from(f: Fallback) -> Self {
174        Composable::Fallback(f)
175    }
176}
177
178impl From<crate::patterns::MapOver> for Composable {
179    fn from(m: crate::patterns::MapOver) -> Self {
180        Composable::MapOver(m)
181    }
182}
183
184// ── Compilation: Composable → TextAgent ──
185
186impl Composable {
187    /// Compile this composable tree into an executable `TextAgent`.
188    ///
189    /// Recursively compiles the tree: pipelines become `SequentialTextAgent`,
190    /// fan-outs become `ParallelTextAgent`, loops become `LoopTextAgent`,
191    /// fallbacks become `FallbackTextAgent`, map-overs become
192    /// `MapOverTextAgent`, and agents compile via `AgentBuilder::build()` —
193    /// whose [`ConfigError`] (an MCP toolset on a text agent) is the only way
194    /// this fails.
195    ///
196    /// ```no_run
197    /// # use gemini_adk_fluent_rs::prelude::*;
198    /// # use std::sync::Arc;
199    /// # async fn run(llm: Arc<dyn BaseLlm>) -> Result<(), AgentError> {
200    /// let pipeline = AgentBuilder::new("writer").instruction("Write a draft")
201    ///     >> AgentBuilder::new("reviewer").instruction("Review and improve");
202    ///
203    /// let agent = pipeline.compile(llm)?;
204    /// let state = State::new();
205    /// let result = agent.run(&state).await?;
206    /// # let _ = result; Ok(())
207    /// # }
208    /// ```
209    pub fn compile(self, llm: Arc<dyn BaseLlm>) -> Result<Arc<dyn TextAgent>, ConfigError> {
210        Ok(match self {
211            Composable::Agent(builder) => builder.build(llm)?,
212
213            Composable::Pipeline(pipeline) => {
214                let children = pipeline
215                    .steps
216                    .into_iter()
217                    .map(|step| step.compile(llm.clone()))
218                    .collect::<Result<Vec<Arc<dyn TextAgent>>, ConfigError>>()?;
219                Arc::new(SequentialTextAgent::new(
220                    pipeline.name.as_deref().unwrap_or("pipeline"),
221                    children,
222                ))
223            }
224
225            Composable::FanOut(fan_out) => {
226                let branches = fan_out
227                    .branches
228                    .into_iter()
229                    .map(|branch| branch.compile(llm.clone()))
230                    .collect::<Result<Vec<Arc<dyn TextAgent>>, ConfigError>>()?;
231                Arc::new(ParallelTextAgent::new(
232                    fan_out.name.as_deref().unwrap_or("fan_out"),
233                    branches,
234                ))
235            }
236
237            Composable::MapOver(map) => {
238                let agent = map.agent.build(llm)?;
239                Arc::new(
240                    gemini_adk_rs::text::MapOverTextAgent::new(
241                        map.name.as_deref().unwrap_or("map_over"),
242                        agent,
243                        map.list_key,
244                    )
245                    .item_key(map.item_key)
246                    .output_key(map.output_key),
247                )
248            }
249
250            Composable::Loop(loop_node) => {
251                let middleware = loop_node.middleware;
252                let body = loop_node.body.compile(llm)?;
253                let mut loop_agent = LoopTextAgent::new(
254                    loop_node.name.as_deref().unwrap_or("loop"),
255                    body,
256                    loop_node.max,
257                );
258
259                if let Some(predicate) = loop_node.until {
260                    loop_agent = loop_agent.until(move |state: &gemini_adk_rs::State| {
261                        // Convert State to serde_json::Value for LoopPredicate compatibility.
262                        let keys = state.keys();
263                        let mut map = serde_json::Map::new();
264                        for key in keys {
265                            if let Some(val) = state.get_raw(&key) {
266                                map.insert(key, val);
267                            }
268                        }
269                        predicate.check(&serde_json::Value::Object(map))
270                    });
271                }
272
273                if !middleware.is_empty() {
274                    loop_agent = loop_agent.with_middleware_chain(chain_from(middleware));
275                }
276
277                Arc::new(loop_agent)
278            }
279
280            Composable::Fallback(fallback) => {
281                let middleware = fallback.middleware;
282                let candidates = fallback
283                    .candidates
284                    .into_iter()
285                    .map(|c| c.compile(llm.clone()))
286                    .collect::<Result<Vec<Arc<dyn TextAgent>>, ConfigError>>()?;
287                let mut agent = FallbackTextAgent::new("fallback", candidates);
288                if !middleware.is_empty() {
289                    agent = agent.with_middleware_chain(chain_from(middleware));
290                }
291                Arc::new(agent)
292            }
293        })
294    }
295}
296
297/// Build a [`MiddlewareChain`] from an ordered list of middleware layers.
298fn chain_from(layers: Vec<Arc<dyn Middleware>>) -> MiddlewareChain {
299    let mut chain = MiddlewareChain::new();
300    for layer in layers {
301        chain.add(layer);
302    }
303    chain
304}
305
306impl Composable {
307    /// Attach middleware to a `Loop` or `Fallback` node — the place where
308    /// combinator-level observers (`M::on_loop`, `M::on_fallback`) live.
309    ///
310    /// For other node kinds (single agent, pipeline, fan-out, map-over) this
311    /// is a no-op: attach `M::` middleware to the agent itself via
312    /// [`AgentBuilder::middleware`](crate::builder::AgentBuilder::middleware) instead.
313    pub fn middleware(self, middleware: impl Into<MiddlewareComposite>) -> Self {
314        match self {
315            Composable::Loop(l) => Composable::Loop(l.middleware(middleware)),
316            Composable::Fallback(f) => Composable::Fallback(f.middleware(middleware)),
317            other => other,
318        }
319    }
320}
321
322// ── Safe variant accessors ──
323//
324// These inspect a `Composable` for a specific shape and return `None` when the
325// variant does not match, rather than panicking. Callers that want the
326// underlying structure should pattern-match directly; these are convenience
327// accessors for introspection (tests, tooling, debugging).
328
329impl Composable {
330    /// The first step of a [`Pipeline`], or `None` if this is not a pipeline
331    /// (or the pipeline is empty).
332    pub fn first_step(&self) -> Option<&Composable> {
333        match self {
334            Composable::Pipeline(p) => p.steps.first(),
335            _ => None,
336        }
337    }
338
339    /// The last step of a [`Pipeline`], or `None` if this is not a pipeline
340    /// (or the pipeline is empty).
341    pub fn last_step(&self) -> Option<&Composable> {
342        match self {
343            Composable::Pipeline(p) => p.steps.last(),
344            _ => None,
345        }
346    }
347
348    /// The `n`th step of a [`Pipeline`], or `None` if this is not a pipeline
349    /// or the index is out of bounds.
350    pub fn nth_step(&self, n: usize) -> Option<&Composable> {
351        match self {
352            Composable::Pipeline(p) => p.steps.get(n),
353            _ => None,
354        }
355    }
356
357    /// All steps of a [`Pipeline`], or `None` if this is not a pipeline.
358    pub fn pipeline_steps(&self) -> Option<&[Composable]> {
359        match self {
360            Composable::Pipeline(p) => Some(&p.steps),
361            _ => None,
362        }
363    }
364
365    /// The branches of a [`FanOut`], or `None` if this is not a fan-out.
366    pub fn fan_out_branches(&self) -> Option<&[Composable]> {
367        match self {
368            Composable::FanOut(f) => Some(&f.branches),
369            _ => None,
370        }
371    }
372
373    /// The termination predicate of a [`Loop`], or `None` if this is not a loop
374    /// (or the loop has no predicate).
375    pub fn loop_predicate(&self) -> Option<&LoopPredicate> {
376        match self {
377            Composable::Loop(l) => l.until.as_ref(),
378            _ => None,
379        }
380    }
381
382    /// The body of a [`Loop`], or `None` if this is not a loop.
383    pub fn loop_body(&self) -> Option<&Composable> {
384        match self {
385            Composable::Loop(l) => Some(&l.body),
386            _ => None,
387        }
388    }
389
390    /// The candidates of a [`Fallback`] chain, or `None` if this is not a fallback.
391    pub fn fallback_candidates(&self) -> Option<&[Composable]> {
392        match self {
393            Composable::Fallback(f) => Some(&f.candidates),
394            _ => None,
395        }
396    }
397}
398
399// ── Pipeline construction helpers ──
400
401impl Pipeline {
402    /// Create a pipeline from the given steps.
403    pub fn new(steps: Vec<Composable>) -> Self {
404        Self {
405            steps,
406            ..Default::default()
407        }
408    }
409
410    /// Create an empty named pipeline (fluent builder entry point). The name
411    /// becomes the compiled `SequentialTextAgent`'s name.
412    ///
413    /// ```
414    /// # use gemini_adk_fluent_rs::prelude::*;
415    /// let etl = Pipeline::builder("etl")
416    ///     .step(AgentBuilder::new("extract"))
417    ///     .step(AgentBuilder::new("transform"))
418    ///     .step(AgentBuilder::new("load"));
419    /// assert_eq!(etl.name.as_deref(), Some("etl"));
420    /// ```
421    pub fn builder(name: &str) -> Self {
422        Self {
423            name: Some(name.to_string()),
424            ..Default::default()
425        }
426    }
427
428    /// Add a sequential step to this pipeline (fluent builder).
429    pub fn step(mut self, agent: impl Into<Composable>) -> Self {
430        self.steps.push(agent.into());
431        self
432    }
433
434    /// Add a sub-agent step (alias for `step` — matches upstream naming).
435    pub fn sub_agent(self, agent: AgentBuilder) -> Self {
436        self.step(agent)
437    }
438
439    /// Set a description (metadata: shown in `Debug` output, not sent to the model).
440    pub fn describe(mut self, desc: &str) -> Self {
441        self.description = Some(desc.to_string());
442        self
443    }
444
445    /// Flatten: if a step is itself a Pipeline, inline its steps.
446    fn push_flat(&mut self, step: Composable) {
447        match step {
448            Composable::Pipeline(p) => self.steps.extend(p.steps),
449            other => self.steps.push(other),
450        }
451    }
452}
453
454impl FanOut {
455    /// Create a fan-out from the given branches.
456    pub fn new(branches: Vec<Composable>) -> Self {
457        Self {
458            branches,
459            ..Default::default()
460        }
461    }
462
463    /// Create an empty named fan-out (fluent builder entry point).
464    ///
465    /// ```
466    /// # use gemini_adk_fluent_rs::prelude::*;
467    /// let research = FanOut::builder("research")
468    ///     .branch(AgentBuilder::new("web"))
469    ///     .branch(AgentBuilder::new("db"));
470    /// assert_eq!(research.branches.len(), 2);
471    /// ```
472    pub fn builder(name: &str) -> Self {
473        Self {
474            name: Some(name.to_string()),
475            ..Default::default()
476        }
477    }
478
479    /// Add a parallel branch (fluent builder).
480    pub fn branch(mut self, agent: impl Into<Composable>) -> Self {
481        self.branches.push(agent.into());
482        self
483    }
484
485    /// Add a sub-agent branch (alias for `branch` — matches upstream naming).
486    pub fn sub_agent(self, agent: AgentBuilder) -> Self {
487        self.branch(agent)
488    }
489
490    /// Set a description (metadata: shown in `Debug` output, not sent to the model).
491    pub fn describe(mut self, desc: &str) -> Self {
492        self.description = Some(desc.to_string());
493        self
494    }
495
496    fn push_flat(&mut self, branch: Composable) {
497        match branch {
498            Composable::FanOut(f) => self.branches.extend(f.branches),
499            other => self.branches.push(other),
500        }
501    }
502}
503
504impl Fallback {
505    /// Create a fallback chain from the given candidates.
506    pub fn new(candidates: Vec<Composable>) -> Self {
507        Self {
508            candidates,
509            middleware: Vec::new(),
510        }
511    }
512
513    /// Attach middleware to the fallback agent (e.g. `M::on_fallback(|name| …)`),
514    /// observed when a fallback branch activates.
515    pub fn middleware(mut self, middleware: impl Into<MiddlewareComposite>) -> Self {
516        self.middleware.extend(middleware.into().layers);
517        self
518    }
519
520    fn push_flat(&mut self, candidate: Composable) {
521        match candidate {
522            Composable::Fallback(f) => self.candidates.extend(f.candidates),
523            other => self.candidates.push(other),
524        }
525    }
526}
527
528// ── Operator: >> (Shr) = Sequential Pipeline ──
529
530/// AgentBuilder >> AgentBuilder → Pipeline
531impl std::ops::Shr for AgentBuilder {
532    type Output = Composable;
533
534    fn shr(self, rhs: AgentBuilder) -> Self::Output {
535        Composable::Pipeline(Pipeline::new(vec![
536            Composable::Agent(self),
537            Composable::Agent(rhs),
538        ]))
539    }
540}
541
542/// Composable >> AgentBuilder → Pipeline (flattening)
543impl std::ops::Shr<AgentBuilder> for Composable {
544    type Output = Composable;
545
546    fn shr(self, rhs: AgentBuilder) -> Self::Output {
547        let mut pipeline = match self {
548            Composable::Pipeline(p) => p,
549            other => Pipeline::new(vec![other]),
550        };
551        pipeline.push_flat(Composable::Agent(rhs));
552        Composable::Pipeline(pipeline)
553    }
554}
555
556/// AgentBuilder >> Composable → Pipeline (flattening)
557impl std::ops::Shr<Composable> for AgentBuilder {
558    type Output = Composable;
559
560    fn shr(self, rhs: Composable) -> Self::Output {
561        let mut pipeline = Pipeline::new(vec![Composable::Agent(self)]);
562        pipeline.push_flat(rhs);
563        Composable::Pipeline(pipeline)
564    }
565}
566
567/// Composable >> Composable → Pipeline (flattening)
568impl std::ops::Shr for Composable {
569    type Output = Composable;
570
571    fn shr(self, rhs: Composable) -> Self::Output {
572        let mut pipeline = match self {
573            Composable::Pipeline(p) => p,
574            other => Pipeline::new(vec![other]),
575        };
576        pipeline.push_flat(rhs);
577        Composable::Pipeline(pipeline)
578    }
579}
580
581// ── Operator: | (BitOr) = Parallel Fan-Out ──
582
583/// AgentBuilder | AgentBuilder → FanOut
584impl std::ops::BitOr for AgentBuilder {
585    type Output = Composable;
586
587    fn bitor(self, rhs: AgentBuilder) -> Self::Output {
588        Composable::FanOut(FanOut::new(vec![
589            Composable::Agent(self),
590            Composable::Agent(rhs),
591        ]))
592    }
593}
594
595/// Composable | AgentBuilder → FanOut (flattening)
596impl std::ops::BitOr<AgentBuilder> for Composable {
597    type Output = Composable;
598
599    fn bitor(self, rhs: AgentBuilder) -> Self::Output {
600        let mut fan_out = match self {
601            Composable::FanOut(f) => f,
602            other => FanOut::new(vec![other]),
603        };
604        fan_out.push_flat(Composable::Agent(rhs));
605        Composable::FanOut(fan_out)
606    }
607}
608
609/// Composable | Composable → FanOut (flattening)
610impl std::ops::BitOr for Composable {
611    type Output = Composable;
612
613    fn bitor(self, rhs: Composable) -> Self::Output {
614        let mut fan_out = match self {
615            Composable::FanOut(f) => f,
616            other => FanOut::new(vec![other]),
617        };
618        fan_out.push_flat(rhs);
619        Composable::FanOut(fan_out)
620    }
621}
622
623// ── Operator: * (Mul<u32>) = Fixed Loop ──
624
625/// AgentBuilder * 3 → Loop(max=3)
626impl std::ops::Mul<u32> for AgentBuilder {
627    type Output = Composable;
628
629    fn mul(self, rhs: u32) -> Self::Output {
630        Composable::Loop(Loop {
631            body: Box::new(Composable::Agent(self)),
632            max: rhs,
633            until: None,
634            middleware: Vec::new(),
635            name: None,
636            description: None,
637        })
638    }
639}
640
641/// Composable * 3 → Loop(max=3)
642impl std::ops::Mul<u32> for Composable {
643    type Output = Composable;
644
645    fn mul(self, rhs: u32) -> Self::Output {
646        Composable::Loop(Loop {
647            body: Box::new(self),
648            max: rhs,
649            until: None,
650            middleware: Vec::new(),
651            name: None,
652            description: None,
653        })
654    }
655}
656
657/// AgentBuilder * until(pred) → conditional Loop
658impl std::ops::Mul<LoopPredicate> for AgentBuilder {
659    type Output = Composable;
660
661    fn mul(self, rhs: LoopPredicate) -> Self::Output {
662        Composable::Loop(Loop {
663            body: Box::new(Composable::Agent(self)),
664            max: u32::MAX,
665            until: Some(rhs),
666            middleware: Vec::new(),
667            name: None,
668            description: None,
669        })
670    }
671}
672
673/// Composable * until(pred) → conditional Loop
674impl std::ops::Mul<LoopPredicate> for Composable {
675    type Output = Composable;
676
677    fn mul(self, rhs: LoopPredicate) -> Self::Output {
678        Composable::Loop(Loop {
679            body: Box::new(self),
680            max: u32::MAX,
681            until: Some(rhs),
682            middleware: Vec::new(),
683            name: None,
684            description: None,
685        })
686    }
687}
688
689// ── Operator: / (Div) = Fallback Chain ──
690// Note: Rust doesn't have a `//` operator. We use `/` (Div) instead.
691
692/// AgentBuilder / AgentBuilder → Fallback
693impl std::ops::Div for AgentBuilder {
694    type Output = Composable;
695
696    fn div(self, rhs: AgentBuilder) -> Self::Output {
697        Composable::Fallback(Fallback::new(vec![
698            Composable::Agent(self),
699            Composable::Agent(rhs),
700        ]))
701    }
702}
703
704/// Composable / AgentBuilder → Fallback (flattening)
705impl std::ops::Div<AgentBuilder> for Composable {
706    type Output = Composable;
707
708    fn div(self, rhs: AgentBuilder) -> Self::Output {
709        let mut fallback = match self {
710            Composable::Fallback(f) => f,
711            other => Fallback::new(vec![other]),
712        };
713        fallback.push_flat(Composable::Agent(rhs));
714        Composable::Fallback(fallback)
715    }
716}
717
718/// Composable / Composable → Fallback (flattening)
719impl std::ops::Div for Composable {
720    type Output = Composable;
721
722    fn div(self, rhs: Composable) -> Self::Output {
723        let mut fallback = match self {
724            Composable::Fallback(f) => f,
725            other => Fallback::new(vec![other]),
726        };
727        fallback.push_flat(rhs);
728        Composable::Fallback(fallback)
729    }
730}
731
732// ── Loop builder method (for chaining max on until-loops) ──
733
734impl Loop {
735    /// Create a loop builder with a body agent and default max iterations.
736    ///
737    /// ```
738    /// # use gemini_adk_fluent_rs::prelude::*;
739    /// let refine = Loop::builder("refine")
740    ///     .step(AgentBuilder::new("refine"))
741    ///     .max_iterations(5);
742    /// assert_eq!(refine.max, 5);
743    /// ```
744    pub fn builder(name: &str) -> Self {
745        Self {
746            body: Box::new(Composable::Pipeline(Pipeline::new(Vec::new()))),
747            max: 10,
748            until: None,
749            middleware: Vec::new(),
750            name: Some(name.to_string()),
751            description: None,
752        }
753    }
754
755    /// Attach middleware to the loop agent (e.g. `M::on_loop(|i| …)`), observed
756    /// on every iteration.
757    pub fn middleware(mut self, middleware: impl Into<MiddlewareComposite>) -> Self {
758        self.middleware.extend(middleware.into().layers);
759        self
760    }
761
762    /// Set the body composable to loop over.
763    pub fn step(mut self, agent: impl Into<Composable>) -> Self {
764        self.body = Box::new(agent.into());
765        self
766    }
767
768    /// Set a maximum number of iterations.
769    pub fn max_iterations(mut self, n: u32) -> Self {
770        self.max = n;
771        self
772    }
773
774    /// Set a description (metadata: shown in `Debug` output, not sent to the model).
775    pub fn describe(mut self, desc: &str) -> Self {
776        self.description = Some(desc.to_string());
777        self
778    }
779}
780
781#[cfg(test)]
782mod tests {
783    use super::*;
784
785    fn agent(name: &str) -> AgentBuilder {
786        AgentBuilder::new(name)
787    }
788
789    #[test]
790    fn pipeline_from_shr() {
791        let result = agent("a") >> agent("b");
792        match result {
793            Composable::Pipeline(p) => assert_eq!(p.steps.len(), 2),
794            _ => panic!("expected Pipeline"),
795        }
796    }
797
798    #[test]
799    fn pipeline_flattens() {
800        let result = agent("a") >> agent("b") >> agent("c");
801        match result {
802            Composable::Pipeline(p) => assert_eq!(p.steps.len(), 3),
803            _ => panic!("expected Pipeline"),
804        }
805    }
806
807    #[test]
808    fn fan_out_from_bitor() {
809        let result = agent("a") | agent("b");
810        match result {
811            Composable::FanOut(f) => assert_eq!(f.branches.len(), 2),
812            _ => panic!("expected FanOut"),
813        }
814    }
815
816    #[test]
817    fn fan_out_flattens() {
818        let result = (agent("a") | agent("b")) | agent("c");
819        match result {
820            Composable::FanOut(f) => assert_eq!(f.branches.len(), 3),
821            _ => panic!("expected FanOut"),
822        }
823    }
824
825    #[test]
826    fn fixed_loop_from_mul() {
827        let result = agent("a") * 3;
828        match result {
829            Composable::Loop(l) => {
830                assert_eq!(l.max, 3);
831                assert!(l.until.is_none());
832            }
833            _ => panic!("expected Loop"),
834        }
835    }
836
837    #[test]
838    fn conditional_loop_from_mul_until() {
839        let pred = until(|_v| true);
840        let result = agent("a") * pred;
841        match result {
842            Composable::Loop(l) => {
843                assert_eq!(l.max, u32::MAX);
844                assert!(l.until.is_some());
845            }
846            _ => panic!("expected Loop"),
847        }
848    }
849
850    #[test]
851    fn fallback_from_div() {
852        let result = agent("a") / agent("b");
853        match result {
854            Composable::Fallback(f) => assert_eq!(f.candidates.len(), 2),
855            _ => panic!("expected Fallback"),
856        }
857    }
858
859    #[test]
860    fn fallback_flattens() {
861        let result = (agent("a") / agent("b")) / agent("c");
862        match result {
863            Composable::Fallback(f) => assert_eq!(f.candidates.len(), 3),
864            _ => panic!("expected Fallback"),
865        }
866    }
867
868    #[test]
869    fn mixed_pipeline_with_fan_out() {
870        let result = agent("a") >> (agent("b") | agent("c"));
871        match &result {
872            Composable::Pipeline(p) => {
873                assert_eq!(p.steps.len(), 2);
874                assert!(matches!(&p.steps[1], Composable::FanOut(_)));
875            }
876            _ => panic!("expected Pipeline"),
877        }
878    }
879
880    #[test]
881    fn pipeline_then_loop() {
882        let result = agent("a") >> (agent("b") * 5);
883        match &result {
884            Composable::Pipeline(p) => {
885                assert_eq!(p.steps.len(), 2);
886                assert!(matches!(&p.steps[1], Composable::Loop(_)));
887            }
888            _ => panic!("expected Pipeline"),
889        }
890    }
891
892    #[test]
893    fn safe_accessors_return_some_on_match() {
894        let pipeline = agent("a").instruction("x") >> agent("b").instruction("y");
895        assert!(pipeline.first_step().is_some());
896        assert!(pipeline.last_step().is_some());
897        assert!(pipeline.nth_step(1).is_some());
898        assert!(pipeline.nth_step(99).is_none());
899        assert_eq!(pipeline.pipeline_steps().map(<[Composable]>::len), Some(2));
900
901        let fan_out = Composable::Agent(agent("a")) | Composable::Agent(agent("b"));
902        assert_eq!(fan_out.fan_out_branches().map(<[Composable]>::len), Some(2));
903
904        let looped = agent("a") * until(|_| true);
905        assert!(looped.loop_predicate().is_some());
906        assert!(looped.loop_body().is_some());
907
908        let fallback = agent("a") / agent("b");
909        assert_eq!(
910            fallback.fallback_candidates().map(<[Composable]>::len),
911            Some(2)
912        );
913    }
914
915    #[test]
916    fn safe_accessors_return_none_on_mismatch() {
917        // Calling a pipeline accessor on a non-Pipeline returns None, not panic.
918        let solo = Composable::Agent(agent("solo"));
919        assert!(solo.first_step().is_none());
920        assert!(solo.last_step().is_none());
921        assert!(solo.nth_step(0).is_none());
922        assert!(solo.pipeline_steps().is_none());
923        assert!(solo.fan_out_branches().is_none());
924        assert!(solo.loop_predicate().is_none());
925        assert!(solo.loop_body().is_none());
926        assert!(solo.fallback_candidates().is_none());
927
928        // A fixed loop (no predicate) returns None for loop_predicate but
929        // Some for loop_body.
930        let fixed = agent("a") * 3;
931        assert!(fixed.loop_predicate().is_none());
932        assert!(fixed.loop_body().is_some());
933        // And a pipeline accessor on a loop is None.
934        assert!(fixed.first_step().is_none());
935    }
936
937    #[test]
938    fn loop_predicate_check() {
939        let pred = until(|v| {
940            v.get("done")
941                .and_then(serde_json::Value::as_bool)
942                .unwrap_or(false)
943        });
944        assert!(!pred.check(&serde_json::json!({"done": false})));
945        assert!(pred.check(&serde_json::json!({"done": true})));
946    }
947
948    // ── compile() tests ──
949
950    mod compile_tests {
951        use super::*;
952        use async_trait::async_trait;
953        use gemini_adk_rs::llm::{BaseLlm, LlmError, LlmRequest, LlmResponse};
954        use gemini_genai_rs::prelude::{Content, Part, Role};
955
956        /// A mock LLM that returns its agent's name from the system instruction.
957        struct NameEchoLlm;
958
959        #[async_trait]
960        impl BaseLlm for NameEchoLlm {
961            fn model_id(&self) -> &str {
962                "name-echo"
963            }
964            async fn generate(&self, req: LlmRequest) -> Result<LlmResponse, LlmError> {
965                let text = req
966                    .system_instruction
967                    .unwrap_or_else(|| "no-instruction".into());
968                Ok(LlmResponse {
969                    content: Content {
970                        role: Some(Role::Model),
971                        parts: vec![Part::Text { text }],
972                    },
973                    finish_reason: Some("STOP".into()),
974                    usage: None,
975                })
976            }
977        }
978
979        fn llm() -> Arc<dyn BaseLlm> {
980            Arc::new(NameEchoLlm)
981        }
982
983        #[tokio::test]
984        async fn compile_single_agent() {
985            let composable = Composable::Agent(AgentBuilder::new("solo").instruction("hello"));
986            let agent = composable.compile(llm()).unwrap();
987            let state = gemini_adk_rs::State::new();
988            let result = agent.run(&state).await.unwrap();
989            assert_eq!(result, "hello");
990        }
991
992        #[tokio::test]
993        async fn compile_pipeline() {
994            let pipeline = agent("a").instruction("step-a") >> agent("b").instruction("step-b");
995            let compiled = pipeline.compile(llm()).unwrap();
996            let state = gemini_adk_rs::State::new();
997            let result = compiled.run(&state).await.unwrap();
998            // Sequential: last agent's output wins. step-b echoes its instruction.
999            assert_eq!(result, "step-b");
1000        }
1001
1002        #[tokio::test]
1003        async fn compile_fan_out() {
1004            let fan_out = Composable::Agent(agent("a").instruction("branch-a"))
1005                | Composable::Agent(agent("b").instruction("branch-b"));
1006            let compiled = fan_out.compile(llm()).unwrap();
1007            let state = gemini_adk_rs::State::new();
1008            let result = compiled.run(&state).await.unwrap();
1009            assert!(result.contains("branch-a"));
1010            assert!(result.contains("branch-b"));
1011        }
1012
1013        #[tokio::test]
1014        async fn compile_loop() {
1015            let looped = agent("counter").instruction("tick") * 3;
1016            let compiled = looped.compile(llm()).unwrap();
1017            let state = gemini_adk_rs::State::new();
1018            let result = compiled.run(&state).await.unwrap();
1019            assert_eq!(result, "tick");
1020        }
1021
1022        #[tokio::test]
1023        async fn compile_fallback() {
1024            let fallback = agent("a").instruction("first") / agent("b").instruction("second");
1025            let compiled = fallback.compile(llm()).unwrap();
1026            let state = gemini_adk_rs::State::new();
1027            let result = compiled.run(&state).await.unwrap();
1028            // First agent succeeds, so its result is returned.
1029            assert_eq!(result, "first");
1030        }
1031
1032        #[tokio::test]
1033        async fn on_loop_fires_through_operator() {
1034            use crate::compose::M;
1035            use std::sync::atomic::{AtomicU32, Ordering};
1036
1037            let count = Arc::new(AtomicU32::new(0));
1038            let c2 = count.clone();
1039            // Attach the combinator-level observer to the loop node.
1040            let looped =
1041                (agent("counter").instruction("tick") * 3).middleware(M::on_loop(move |_i| {
1042                    c2.fetch_add(1, Ordering::SeqCst);
1043                }));
1044            let compiled = looped.compile(llm()).unwrap();
1045            let state = gemini_adk_rs::State::new();
1046            compiled.run(&state).await.unwrap();
1047            // Three iterations → three LoopIteration events observed.
1048            assert_eq!(count.load(Ordering::SeqCst), 3);
1049        }
1050
1051        #[tokio::test]
1052        async fn compile_loop_with_predicate() {
1053            // Use a mock LLM that increments state on each call.
1054            struct IncrementLlm;
1055
1056            #[async_trait]
1057            impl BaseLlm for IncrementLlm {
1058                fn model_id(&self) -> &str {
1059                    "incr"
1060                }
1061                async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
1062                    Ok(LlmResponse {
1063                        content: Content {
1064                            role: Some(Role::Model),
1065                            parts: vec![Part::Text {
1066                                text: "done".into(),
1067                            }],
1068                        },
1069                        finish_reason: Some("STOP".into()),
1070                        usage: None,
1071                    })
1072                }
1073            }
1074
1075            // Build a FnTextAgent-driven loop instead to test predicate.
1076            // We'll test via the operators directly.
1077            let pred = until(|v| v.get("n").and_then(serde_json::Value::as_i64).unwrap_or(0) >= 3);
1078            let body = agent("incr").instruction("increment");
1079            let looped = body * pred;
1080
1081            // Compile it. The predicate checks state for "n" >= 3, but
1082            // the mock LLM doesn't set "n". Loop will run max iterations.
1083            // This tests that the predicate is wired through.
1084            let compiled = looped.compile(Arc::new(IncrementLlm)).unwrap();
1085            let state = gemini_adk_rs::State::new();
1086            let _ = state.set("n", 5); // Pre-set to pass predicate immediately.
1087            let result = compiled.run(&state).await.unwrap();
1088            assert_eq!(result, "done"); // Ran once, predicate passed.
1089        }
1090
1091        #[tokio::test]
1092        async fn compile_mixed_pipeline_with_fan_out() {
1093            let mixed = agent("a").instruction("start")
1094                >> (Composable::Agent(agent("b").instruction("left"))
1095                    | Composable::Agent(agent("c").instruction("right")));
1096            let compiled = mixed.compile(llm()).unwrap();
1097            let state = gemini_adk_rs::State::new();
1098            let result = compiled.run(&state).await.unwrap();
1099            assert!(result.contains("left"));
1100            assert!(result.contains("right"));
1101        }
1102    }
1103}
gemini_adk_fluent_rs/operators.rs

gemini_adk_fluent_rs/
operators.rs
