gemini_adk_fluent_rs/
patterns.rs

1//! Pre-built patterns — common multi-agent workflows.
2//!
3//! High-level functions that compose agents into standard patterns:
4//! review loops, cascades, fan-out-merge, supervised workflows, etc.
5//!
6//! Each function returns a [`Composable`] that can be compiled into an
7//! executable [`TextAgent`](gemini_adk_rs::text::TextAgent) via
8//! [`Composable::compile()`](crate::operators::Composable::compile).
9//!
10//! # Examples
11//!
12//! ```
13//! use gemini_adk_fluent_rs::prelude::*;
14//!
15//! // Review loop: author writes, reviewer checks, loop until approved
16//! let draft = review_loop(
17//!     AgentBuilder::new("author").instruction("Write an essay"),
18//!     AgentBuilder::new("reviewer").instruction("Review and set approved=true when good"),
19//!     3,
20//! );
21//!
22//! // Cascade: try agents in order, first success wins
23//! let robust = cascade(vec![
24//!     AgentBuilder::new("primary"),
25//!     AgentBuilder::new("fallback"),
26//! ]);
27//!
28//! // Fan-out-merge: parallel agents, then merge
29//! let research = fan_out_merge(
30//!     vec![AgentBuilder::new("web"), AgentBuilder::new("db")],
31//!     AgentBuilder::new("synthesizer"),
32//! );
33//! # let _ = (draft, robust, research);
34//! ```
35
36use crate::builder::AgentBuilder;
37use crate::operators::{Branch, Composable, Fallback, FanOut, Loop, LoopPredicate, Pipeline};
38
39/// Review loop: author writes, reviewer checks, loop until approved.
40///
41/// The author agent produces output, then the reviewer evaluates it.
42/// The loop terminates when the reviewer sets `"approved"` to `true`
43/// in the state, or after `max_rounds` iterations.
44///
45/// # Arguments
46///
47/// * `author` — The agent that produces drafts.
48/// * `reviewer` — The agent that evaluates and sets `"approved": true` when satisfied.
49/// * `max_rounds` — Maximum number of author-reviewer cycles.
50///
51/// # Example
52///
53/// ```no_run
54/// # use gemini_adk_fluent_rs::prelude::*;
55/// # use std::sync::Arc;
56/// # fn run(llm: Arc<dyn BaseLlm>) -> Result<(), ConfigError> {
57/// let workflow = review_loop(
58///     AgentBuilder::new("writer").instruction("Write a blog post"),
59///     AgentBuilder::new("editor").instruction("Review. Set approved=true if publication-ready."),
60///     3,
61/// );
62/// let agent = workflow.compile(llm)?;
63/// # let _ = agent; Ok(())
64/// # }
65/// ```
66pub fn review_loop(author: AgentBuilder, reviewer: AgentBuilder, max_rounds: usize) -> Composable {
67    let inner = Composable::Pipeline(Pipeline::new(vec![
68        Composable::Agent(author),
69        Composable::Agent(reviewer),
70    ]));
71
72    Composable::Loop(Loop {
73        body: Box::new(inner),
74        max: max_rounds as u32,
75        middleware: Vec::new(),
76        name: None,
77        description: None,
78        until: Some(LoopPredicate::new(|state| {
79            state
80                .get("approved")
81                .and_then(serde_json::Value::as_bool)
82                .unwrap_or(false)
83        })),
84    })
85}
86
87/// Review loop with a custom quality key and target value.
88///
89/// Like [`review_loop`] but allows specifying which state key the reviewer
90/// writes to and what value signals completion.
91///
92/// # Arguments
93///
94/// * `worker` — The agent that produces output.
95/// * `reviewer` — The agent that evaluates quality.
96/// * `quality_key` — State key the reviewer writes (e.g., `"quality"`).
97/// * `target` — Value of `quality_key` that signals completion (e.g., `"good"`).
98/// * `max_rounds` — Maximum iterations.
99pub fn review_loop_keyed(
100    worker: AgentBuilder,
101    reviewer: AgentBuilder,
102    quality_key: &str,
103    target: &str,
104    max_rounds: u32,
105) -> Composable {
106    let key = quality_key.to_string();
107    let target = target.to_string();
108
109    let inner = Composable::Pipeline(Pipeline::new(vec![
110        Composable::Agent(worker),
111        Composable::Agent(reviewer),
112    ]));
113
114    Composable::Loop(Loop {
115        body: Box::new(inner),
116        max: max_rounds,
117        middleware: Vec::new(),
118        name: None,
119        description: None,
120        until: Some(LoopPredicate::new(move |state| {
121            state
122                .get(&key)
123                .and_then(|v| v.as_str())
124                .map(|v| v == target)
125                .unwrap_or(false)
126        })),
127    })
128}
129
130/// Cascade: try agents in sequence, first success wins.
131///
132/// This is an alias for a fallback chain. Each agent is tried in order;
133/// the first one that succeeds provides the result.
134///
135/// # Example
136///
137/// ```
138/// # use gemini_adk_fluent_rs::prelude::*;
139/// let robust = cascade(vec![
140///     AgentBuilder::new("fast").instruction("Quick answer"),
141///     AgentBuilder::new("thorough").instruction("Detailed answer"),
142/// ]);
143/// assert!(matches!(robust, Composable::Fallback(_)));
144/// ```
145pub fn cascade(agents: Vec<AgentBuilder>) -> Composable {
146    Composable::Fallback(Fallback::new(
147        agents.into_iter().map(Composable::Agent).collect(),
148    ))
149}
150
151/// Fan-out-merge: run agents in parallel, then merge results with a merger agent.
152///
153/// All `agents` execute concurrently via fan-out. Their combined output is
154/// then fed into the `merger` agent, which synthesizes a final result.
155///
156/// # Arguments
157///
158/// * `agents` — Agents to run in parallel.
159/// * `merger` — Agent that merges the parallel results.
160///
161/// # Example
162///
163/// ```
164/// # use gemini_adk_fluent_rs::prelude::*;
165/// let research = fan_out_merge(
166///     vec![
167///         AgentBuilder::new("web-search").instruction("Search the web"),
168///         AgentBuilder::new("db-lookup").instruction("Query the database"),
169///     ],
170///     AgentBuilder::new("synthesizer").instruction("Combine research findings"),
171/// );
172/// assert!(matches!(research, Composable::Pipeline(_)));
173/// ```
174pub fn fan_out_merge(agents: Vec<AgentBuilder>, merger: AgentBuilder) -> Composable {
175    let fan_out = Composable::FanOut(FanOut::new(
176        agents.into_iter().map(Composable::Agent).collect(),
177    ));
178
179    Composable::Pipeline(Pipeline::new(vec![fan_out, Composable::Agent(merger)]))
180}
181
182/// Chain: simple sequential pipeline of agents.
183///
184/// This is an alias for the `>>` operator but accepts a `Vec`.
185/// Each agent runs in order, with the output of one feeding into the next.
186///
187/// # Example
188///
189/// ```
190/// # use gemini_adk_fluent_rs::prelude::*;
191/// let pipeline = chain(vec![
192///     AgentBuilder::new("extract"),
193///     AgentBuilder::new("transform"),
194///     AgentBuilder::new("load"),
195/// ]);
196/// assert!(matches!(pipeline, Composable::Pipeline(_)));
197/// ```
198pub fn chain(agents: Vec<AgentBuilder>) -> Composable {
199    Composable::Pipeline(Pipeline::new(
200        agents.into_iter().map(Composable::Agent).collect(),
201    ))
202}
203
204/// Conditional: run one of two workflows, chosen by a state predicate.
205///
206/// `predicate` sees the state (as a JSON object) when the compiled agent
207/// runs. If it returns `true`, `if_true` runs; otherwise `if_false` runs.
208/// Exactly one branch runs, and each branch keeps its full configuration.
209///
210/// # Arguments
211///
212/// * `predicate` — Function that inspects state (as `serde_json::Value`) and returns a bool.
213/// * `if_true` — Agent or workflow to run when the predicate is true.
214/// * `if_false` — Agent or workflow to run when the predicate is false.
215///
216/// # Example
217///
218/// ```
219/// # use gemini_adk_fluent_rs::prelude::*;
220/// let routed = conditional(
221///     |state| state.get("premium").and_then(|v| v.as_bool()).unwrap_or(false),
222///     AgentBuilder::new("premium-agent").instruction("Full-featured response"),
223///     AgentBuilder::new("basic-agent").instruction("Basic response"),
224/// );
225/// assert!(matches!(routed, Composable::Branch(_)));
226/// ```
227pub fn conditional(
228    predicate: impl Fn(&serde_json::Value) -> bool + Send + Sync + 'static,
229    if_true: impl Into<Composable>,
230    if_false: impl Into<Composable>,
231) -> Composable {
232    Composable::Branch(Branch {
233        predicate: LoopPredicate::new(predicate),
234        if_true: Box::new(if_true.into()),
235        if_false: Box::new(if_false.into()),
236        name: None,
237    })
238}
239
240/// Supervised: worker with supervisor oversight loop.
241///
242/// The worker agent produces output, then the supervisor reviews it.
243/// The loop repeats until the supervisor sets `"approved"` to `true`
244/// in the state, or after `max_rounds` iterations.
245///
246/// This is semantically similar to [`review_loop`] but framed as a
247/// worker-supervisor relationship rather than author-reviewer.
248///
249/// # Arguments
250///
251/// * `worker` — The agent that performs the task.
252/// * `supervisor` — The agent that oversees and approves work.
253/// * `max_rounds` — Maximum number of worker-supervisor cycles.
254///
255/// # Example
256///
257/// ```
258/// # use gemini_adk_fluent_rs::prelude::*;
259/// let managed = supervised(
260///     AgentBuilder::new("coder").instruction("Write the implementation"),
261///     AgentBuilder::new("lead").instruction("Code review. Set approved=true if ready to merge."),
262///     5,
263/// );
264/// assert!(matches!(managed, Composable::Loop(_)));
265/// ```
266pub fn supervised(worker: AgentBuilder, supervisor: AgentBuilder, max_rounds: usize) -> Composable {
267    let inner = Composable::Pipeline(Pipeline::new(vec![
268        Composable::Agent(worker),
269        Composable::Agent(supervisor),
270    ]));
271
272    Composable::Loop(Loop {
273        body: Box::new(inner),
274        max: max_rounds as u32,
275        middleware: Vec::new(),
276        name: None,
277        description: None,
278        until: Some(LoopPredicate::new(|state| {
279            state
280                .get("approved")
281                .and_then(serde_json::Value::as_bool)
282                .unwrap_or(false)
283        })),
284    })
285}
286
287/// Supervised with a custom approval key.
288///
289/// Like [`supervised`] but allows specifying which state key signals approval.
290///
291/// # Arguments
292///
293/// * `worker` — The agent that performs the task.
294/// * `supervisor` — The agent that oversees work.
295/// * `approval_key` — State key the supervisor sets to `true` when satisfied.
296/// * `max_revisions` — Maximum iterations.
297pub fn supervised_keyed(
298    worker: AgentBuilder,
299    supervisor: AgentBuilder,
300    approval_key: &str,
301    max_revisions: u32,
302) -> Composable {
303    let key = approval_key.to_string();
304
305    let inner = Composable::Pipeline(Pipeline::new(vec![
306        Composable::Agent(worker),
307        Composable::Agent(supervisor),
308    ]));
309
310    Composable::Loop(Loop {
311        body: Box::new(inner),
312        max: max_revisions,
313        middleware: Vec::new(),
314        name: None,
315        description: None,
316        until: Some(LoopPredicate::new(move |state| {
317            state
318                .get(&key)
319                .and_then(serde_json::Value::as_bool)
320                .unwrap_or(false)
321        })),
322    })
323}
324
325/// Map-over: apply one agent to every item of a state list.
326///
327/// At run time the compiled node reads the JSON array at `state[list_key]`,
328/// runs `agent` once per item with the item in `state["_item"]` (and as the
329/// agent's `input`), and collects the outputs into `state["_results"]`.
330/// Items run sequentially — `MapOver::item_key`/`output_key` rename the
331/// slots. Composes like any other node:
332///
333/// ```
334/// # use gemini_adk_fluent_rs::prelude::*;
335/// let batch = map_over(AgentBuilder::new("summarize"), "documents")
336///     >> AgentBuilder::new("merge");
337/// assert!(matches!(batch, Composable::Pipeline(_)));
338/// ```
339pub fn map_over(agent: AgentBuilder, list_key: impl Into<String>) -> Composable {
340    Composable::MapOver(MapOver::new(agent, list_key))
341}
342
343/// A map-over workflow node — applies one agent to many items. Build with
344/// [`map_over`]; compiles to `MapOverTextAgent`.
345#[derive(Clone, Debug)]
346pub struct MapOver {
347    /// The agent template applied to each item.
348    pub agent: AgentBuilder,
349    /// State key holding the JSON array to iterate.
350    pub list_key: String,
351    /// State key the current item is written to (default `"_item"`).
352    pub item_key: String,
353    /// State key the collected outputs are written to (default `"_results"`).
354    pub output_key: String,
355    /// Name given to the compiled agent (default `"map_over"`).
356    pub name: Option<String>,
357}
358
359impl MapOver {
360    /// Create a map-over node for `agent` over the list at `list_key`.
361    pub fn new(agent: AgentBuilder, list_key: impl Into<String>) -> Self {
362        Self {
363            agent,
364            list_key: list_key.into(),
365            item_key: "_item".into(),
366            output_key: "_results".into(),
367            name: None,
368        }
369    }
370
371    /// State key the current item is written to for each run.
372    pub fn item_key(mut self, key: impl Into<String>) -> Self {
373        self.item_key = key.into();
374        self
375    }
376
377    /// State key the collected outputs are written to.
378    pub fn output_key(mut self, key: impl Into<String>) -> Self {
379        self.output_key = key.into();
380        self
381    }
382
383    /// Name the compiled agent.
384    pub fn name(mut self, name: impl Into<String>) -> Self {
385        self.name = Some(name.into());
386        self
387    }
388}
389
390/// Map-reduce: map `mapper` over the list at `list_key`, then run `reducer`
391/// over the collected results (`state["_results"]`). A pipeline of a
392/// [`map_over`] node and the reducer.
393pub fn map_reduce(
394    mapper: AgentBuilder,
395    reducer: AgentBuilder,
396    list_key: impl Into<String>,
397) -> Composable {
398    Composable::Pipeline(Pipeline::new(vec![
399        map_over(mapper, list_key),
400        Composable::Agent(reducer),
401    ]))
402}
403
404#[cfg(test)]
405mod tests {
406    use super::*;
407
408    fn agent(name: &str) -> AgentBuilder {
409        AgentBuilder::new(name)
410    }
411
412    #[test]
413    fn review_loop_creates_loop_with_pipeline() {
414        let result = review_loop(agent("writer"), agent("reviewer"), 3);
415        match &result {
416            Composable::Loop(l) => {
417                assert_eq!(l.max, 3);
418                assert!(l.until.is_some());
419                assert!(matches!(&*l.body, Composable::Pipeline(p) if p.steps.len() == 2));
420            }
421            _ => panic!("expected Loop"),
422        }
423    }
424
425    #[test]
426    fn review_loop_predicate_checks_approved() {
427        let result = review_loop(agent("w"), agent("r"), 3);
428        if let Composable::Loop(l) = result {
429            let pred = l.until.unwrap();
430            assert!(!pred.check(&serde_json::json!({"approved": false})));
431            assert!(pred.check(&serde_json::json!({"approved": true})));
432            assert!(!pred.check(&serde_json::json!({})));
433        }
434    }
435
436    #[test]
437    fn review_loop_keyed_predicate_works() {
438        let result = review_loop_keyed(agent("w"), agent("r"), "quality", "good", 3);
439        if let Composable::Loop(l) = result {
440            let pred = l.until.unwrap();
441            assert!(!pred.check(&serde_json::json!({"quality": "bad"})));
442            assert!(pred.check(&serde_json::json!({"quality": "good"})));
443        }
444    }
445
446    #[test]
447    fn cascade_creates_fallback() {
448        let result = cascade(vec![agent("a"), agent("b"), agent("c")]);
449        match result {
450            Composable::Fallback(f) => assert_eq!(f.candidates.len(), 3),
451            _ => panic!("expected Fallback"),
452        }
453    }
454
455    #[test]
456    fn fan_out_merge_creates_pipeline_with_fan_out_then_merger() {
457        let result = fan_out_merge(vec![agent("a"), agent("b")], agent("merger"));
458        match &result {
459            Composable::Pipeline(p) => {
460                assert_eq!(p.steps.len(), 2);
461                assert!(matches!(&p.steps[0], Composable::FanOut(f) if f.branches.len() == 2));
462                assert!(matches!(&p.steps[1], Composable::Agent(a) if a.name() == "merger"));
463            }
464            _ => panic!("expected Pipeline"),
465        }
466    }
467
468    #[test]
469    fn chain_creates_pipeline() {
470        let result = chain(vec![agent("a"), agent("b"), agent("c")]);
471        match result {
472            Composable::Pipeline(p) => assert_eq!(p.steps.len(), 3),
473            _ => panic!("expected Pipeline"),
474        }
475    }
476
477    /// Exactly one branch runs, chosen by state at run time. The old
478    /// implementation ran the true branch whatever the predicate said, and
479    /// rebuilt each branch from its name and instruction alone.
480    #[tokio::test]
481    async fn conditional_runs_exactly_the_chosen_branch() {
482        use gemini_adk_rs::llm::{LlmResponse, MockLlm};
483        use std::sync::Arc;
484
485        // The model answers with the instruction it was given, so the reply
486        // names the branch that ran.
487        let llm = MockLlm::from_fn(|req| {
488            Ok(LlmResponse::from_text(
489                req.system_instruction.clone().unwrap_or_default(),
490            ))
491        });
492        let routed = conditional(
493            |state| {
494                state
495                    .get("flag")
496                    .and_then(serde_json::Value::as_bool)
497                    .unwrap_or(false)
498            },
499            agent("yes").instruction("true branch").temperature(0.3),
500            agent("no").instruction("false branch"),
501        )
502        .compile(Arc::new(llm.clone()))
503        .unwrap();
504
505        let state = gemini_adk_rs::State::new();
506        state.set("flag", true).unwrap();
507        assert_eq!(routed.run(&state).await.unwrap(), "true branch");
508        assert_eq!(llm.last_request().unwrap().temperature, Some(0.3));
509
510        state.set("flag", false).unwrap();
511        assert_eq!(routed.run(&state).await.unwrap(), "false branch");
512        assert_eq!(llm.call_count(), 2, "one model call per run");
513    }
514
515    #[test]
516    fn supervised_creates_loop() {
517        let result = supervised(agent("worker"), agent("supervisor"), 5);
518        match &result {
519            Composable::Loop(l) => {
520                assert_eq!(l.max, 5);
521                assert!(l.until.is_some());
522                assert!(matches!(&*l.body, Composable::Pipeline(p) if p.steps.len() == 2));
523            }
524            _ => panic!("expected Loop"),
525        }
526    }
527
528    #[test]
529    fn supervised_predicate_checks_approved() {
530        let result = supervised(agent("w"), agent("s"), 5);
531        if let Composable::Loop(l) = result {
532            let pred = l.until.unwrap();
533            assert!(!pred.check(&serde_json::json!({"approved": false})));
534            assert!(pred.check(&serde_json::json!({"approved": true})));
535        }
536    }
537
538    #[test]
539    fn supervised_keyed_predicate_works() {
540        let result = supervised_keyed(agent("w"), agent("s"), "approved", 5);
541        if let Composable::Loop(l) = result {
542            let pred = l.until.unwrap();
543            assert!(!pred.check(&serde_json::json!({"approved": false})));
544            assert!(pred.check(&serde_json::json!({"approved": true})));
545        }
546    }
547
548    #[test]
549    fn map_over_is_a_composable_node() {
550        match map_over(agent("processor"), "items") {
551            Composable::MapOver(m) => {
552                assert_eq!(m.agent.name(), "processor");
553                assert_eq!(m.list_key, "items");
554                assert_eq!(m.item_key, "_item");
555            }
556            other => panic!("expected MapOver, got {other:?}"),
557        }
558    }
559
560    #[test]
561    fn map_reduce_is_map_over_then_reducer() {
562        match map_reduce(agent("mapper"), agent("reducer"), "items") {
563            Composable::Pipeline(p) => {
564                assert_eq!(p.steps.len(), 2);
565                assert!(
566                    matches!(&p.steps[0], Composable::MapOver(m) if m.agent.name() == "mapper")
567                );
568                assert!(matches!(&p.steps[1], Composable::Agent(a) if a.name() == "reducer"));
569            }
570            other => panic!("expected Pipeline, got {other:?}"),
571        }
572    }
573
574    #[tokio::test]
575    async fn map_over_compiles_and_runs_per_item() {
576        use gemini_adk_rs::llm::{BaseLlm, LlmError, LlmRequest, LlmResponse};
577        use gemini_genai_rs::prelude::{Content, Part, Role};
578        use std::sync::Arc;
579
580        struct Echo;
581        #[async_trait::async_trait]
582        impl BaseLlm for Echo {
583            fn model_id(&self) -> &str {
584                "echo"
585            }
586            async fn generate(&self, req: LlmRequest) -> Result<LlmResponse, LlmError> {
587                let text = req
588                    .contents
589                    .iter()
590                    .flat_map(|c| &c.parts)
591                    .filter_map(|p| match p {
592                        Part::Text { text } => Some(text.clone()),
593                        _ => None,
594                    })
595                    .collect::<Vec<_>>()
596                    .join("");
597                Ok(LlmResponse {
598                    content: Content {
599                        role: Some(Role::Model),
600                        parts: vec![Part::Text { text }],
601                    },
602                    finish_reason: Some("STOP".into()),
603                    usage: None,
604                })
605            }
606        }
607
608        let node = map_over(agent("echo"), "items");
609        let compiled = node.compile(Arc::new(Echo)).expect("compiles");
610        let state = gemini_adk_rs::State::new();
611        let _ = state.set("items", serde_json::json!(["a", "b"]));
612        let out = compiled.run(&state).await.expect("runs");
613        assert!(out.contains("\"a\"") && out.contains("\"b\""), "{out}");
614        let results: Vec<String> = state.get("_results").unwrap_or_default();
615        assert_eq!(results.len(), 2);
616    }
617}