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

gemini_adk_fluent_rs/
operators.rs
