adk_graph/
agent.rs

1//! GraphAgent - ADK Agent integration for graph workflows
2//!
3//! Provides a builder pattern similar to LlmAgent and RealtimeAgent.
4
5use crate::checkpoint::Checkpointer;
6use crate::deferred::DeferredNodeConfig;
7use crate::edge::{END, Edge, EdgeTarget, START};
8use crate::error::{GraphError, Result};
9use crate::graph::{CompiledGraph, StateGraph};
10use crate::node::{ExecutionConfig, FunctionNode, Node, NodeContext, NodeOutput};
11use crate::state::{State, StateSchema};
12use crate::stream::{StreamEvent, StreamMode};
13use crate::timeout::TimeoutPolicy;
14use adk_core::{
15    Agent, AgentCapabilities, AgentRelationshipKind, AgentTopology, AgentTopologyMember,
16    AgentTopologyRelationship, Content, Event, EventStream, InvocationContext,
17};
18use async_trait::async_trait;
19use serde_json::json;
20use std::collections::HashMap;
21use std::future::Future;
22use std::pin::Pin;
23use std::sync::Arc;
24
25/// Type alias for callbacks
26pub type BeforeAgentCallback = Arc<
27    dyn Fn(Arc<dyn InvocationContext>) -> Pin<Box<dyn Future<Output = adk_core::Result<()>> + Send>>
28        + Send
29        + Sync,
30>;
31
32pub type AfterAgentCallback = Arc<
33    dyn Fn(
34            Arc<dyn InvocationContext>,
35            Event,
36        ) -> Pin<Box<dyn Future<Output = adk_core::Result<()>> + Send>>
37        + Send
38        + Sync,
39>;
40
41/// Type alias for input mapper function
42pub type InputMapper = Arc<dyn Fn(&dyn InvocationContext) -> State + Send + Sync>;
43
44/// Type alias for output mapper function
45pub type OutputMapper = Arc<dyn Fn(&State) -> Vec<Event> + Send + Sync>;
46
47/// GraphAgent wraps a CompiledGraph as an ADK Agent
48pub struct GraphAgent {
49    name: String,
50    description: String,
51    graph: Arc<CompiledGraph>,
52    /// Map InvocationContext to graph input state
53    input_mapper: InputMapper,
54    /// Map graph output state to ADK Events
55    output_mapper: OutputMapper,
56    /// Before agent callback
57    before_callback: Option<BeforeAgentCallback>,
58    /// After agent callback
59    after_callback: Option<AfterAgentCallback>,
60}
61
62impl GraphAgent {
63    /// Create a new GraphAgent builder
64    pub fn builder(name: &str) -> GraphAgentBuilder {
65        GraphAgentBuilder::new(name)
66    }
67
68    /// Create directly from a compiled graph
69    pub fn from_graph(name: &str, graph: CompiledGraph) -> Self {
70        Self {
71            name: name.to_string(),
72            description: String::new(),
73            graph: Arc::new(graph),
74            input_mapper: Arc::new(default_input_mapper),
75            output_mapper: Arc::new(default_output_mapper),
76            before_callback: None,
77            after_callback: None,
78        }
79    }
80
81    /// Build a `GraphAgent` from a `WorkflowSchema`.
82    ///
83    /// Delegates to `schema.build_graph()` to construct the graph from the
84    /// workflow schema's action nodes, edges, and conditions.
85    #[cfg(feature = "action")]
86    pub fn from_workflow_schema(
87        name: &str,
88        schema: &crate::workflow::WorkflowSchema,
89    ) -> Result<Self> {
90        schema.build_graph(name)
91    }
92
93    /// Get the underlying compiled graph
94    pub fn graph(&self) -> &CompiledGraph {
95        &self.graph
96    }
97
98    /// Execute the graph directly (bypassing Agent trait)
99    pub async fn invoke(&self, input: State, config: ExecutionConfig) -> Result<State> {
100        self.graph.invoke(input, config).await
101    }
102
103    /// Stream execution
104    pub fn stream(
105        &self,
106        input: State,
107        config: ExecutionConfig,
108        mode: StreamMode,
109    ) -> impl futures::Stream<Item = Result<StreamEvent>> + '_ {
110        self.graph.stream(input, config, mode)
111    }
112}
113
114#[async_trait]
115impl Agent for GraphAgent {
116    fn name(&self) -> &str {
117        &self.name
118    }
119
120    fn description(&self) -> &str {
121        &self.description
122    }
123
124    fn sub_agents(&self) -> &[Arc<dyn Agent>] {
125        &[]
126    }
127
128    fn supports_agent_transfer(&self) -> bool {
129        false
130    }
131
132    fn capabilities(&self) -> AgentCapabilities {
133        AgentCapabilities {
134            checkpoint_resume: self.graph.has_checkpointer(),
135            shared_state: true,
136            invocation_metadata: true,
137            ..AgentCapabilities::default()
138        }
139    }
140
141    fn topology(&self) -> Option<AgentTopology> {
142        let mut nodes = self.graph.nodes.values().collect::<Vec<_>>();
143        nodes.sort_by_key(|node| node.name().to_string());
144        let mut members = vec![AgentTopologyMember {
145            name: self.name.clone(),
146            description: self.description.clone(),
147            coordinator: true,
148            capabilities: self.capabilities(),
149        }];
150        members.extend(nodes.into_iter().map(|node| AgentTopologyMember {
151            name: node.name().to_string(),
152            description: node.description().to_string(),
153            coordinator: false,
154            capabilities: node.capabilities(),
155        }));
156
157        let mut relationships = self
158            .graph
159            .get_entry_nodes()
160            .into_iter()
161            .map(|entry| AgentTopologyRelationship {
162                from: self.name.clone(),
163                to: entry,
164                kind: AgentRelationshipKind::Flow,
165            })
166            .collect::<Vec<_>>();
167        for edge in &self.graph.edges {
168            match edge {
169                Edge::Direct { source, target: EdgeTarget::Node(target) } => {
170                    relationships.push(AgentTopologyRelationship {
171                        from: source.clone(),
172                        to: target.clone(),
173                        kind: AgentRelationshipKind::Flow,
174                    });
175                }
176                Edge::Conditional { source, targets, .. } => {
177                    relationships.extend(targets.values().filter_map(|target| {
178                        target.node_name().map(|target| AgentTopologyRelationship {
179                            from: source.clone(),
180                            to: target.to_string(),
181                            kind: AgentRelationshipKind::Flow,
182                        })
183                    }));
184                }
185                Edge::Direct { target: EdgeTarget::End, .. } | Edge::Entry { .. } => {}
186            }
187        }
188        relationships.sort_by(|left, right| (&left.from, &left.to).cmp(&(&right.from, &right.to)));
189        relationships.dedup_by(|left, right| left.from == right.from && left.to == right.to);
190
191        Some(AgentTopology {
192            root: self.name.clone(),
193            coordinator: self.name.clone(),
194            members,
195            relationships,
196        })
197    }
198
199    async fn run(&self, ctx: Arc<dyn InvocationContext>) -> adk_core::Result<EventStream> {
200        // Call before callback
201        if let Some(callback) = &self.before_callback {
202            callback(ctx.clone()).await?;
203        }
204
205        // Map context to input state
206        let input = (self.input_mapper)(ctx.as_ref());
207
208        // Create execution config from context, carrying the invocation so an
209        // `AgentNode` inside the graph presents this run's identity, services, and
210        // cancellation rather than a synthetic standalone context.
211        let config = ExecutionConfig::new(ctx.session_id()).with_parent_context(ctx.clone());
212
213        // Execute graph
214        let graph = self.graph.clone();
215        let output_mapper = self.output_mapper.clone();
216        let after_callback = self.after_callback.clone();
217        let ctx_clone = ctx.clone();
218
219        let stream = async_stream::stream! {
220            match graph.invoke(input, config).await {
221                Ok(state) => {
222                    let events = output_mapper(&state);
223                    for event in events {
224                        // Call after callback for each event
225                        if let Some(callback) = &after_callback
226                            && let Err(e) = callback(ctx_clone.clone(), event.clone()).await {
227                                yield Err(e);
228                                return;
229                            }
230                        yield Ok(event);
231                    }
232                }
233                Err(GraphError::Interrupted(interrupt)) => {
234                    // The `Agent` trait yields events, so an interrupt cannot be
235                    // returned as an error without ending the invocation. Emit one
236                    // event carrying the structured pause so a caller can read the
237                    // node, the payload, and the checkpoint to resume from.
238                    let tool_confirmation = crate::interrupt::GraphToolConfirmationPause::from_interrupted_execution(&interrupt);
239                    let payload = tool_confirmation.clone().map_or_else(
240                        || crate::interrupt::GraphInterruptPayload::new(
241                            &interrupt.interrupt,
242                            &interrupt.thread_id,
243                            &interrupt.checkpoint_id,
244                        ),
245                        crate::interrupt::GraphInterruptPayload::from_tool_confirmation_pause,
246                    );
247                    let mut event = Event::new("graph_interrupted");
248                    if let Some(pause) = tool_confirmation {
249                        event.llm_response.interrupted = true;
250                        event.llm_response.turn_complete = true;
251                        event.actions.tool_confirmation = Some(pause.request);
252                    }
253                    event.set_content(
254                        Content::new("assistant").with_text(interrupt.interrupt.to_string()),
255                    );
256                    event.provider_metadata.insert(
257                        crate::interrupt::INTERRUPT_METADATA_KEY.to_string(),
258                        payload.to_metadata_value(),
259                    );
260                    yield Ok(event);
261                }
262                Err(e) => {
263                    yield Err(adk_core::AdkError::agent(e.to_string()));
264                }
265            }
266        };
267
268        Ok(Box::pin(stream))
269    }
270}
271
272/// Default input mapper - extracts content from InvocationContext
273fn default_input_mapper(ctx: &dyn InvocationContext) -> State {
274    let mut state = State::new();
275
276    // Get user content
277    let content = ctx.user_content();
278    let text: String = content.parts.iter().filter_map(|p| p.text()).collect::<Vec<_>>().join("\n");
279
280    if !text.is_empty() {
281        state.insert("input".to_string(), json!(text));
282        state.insert("messages".to_string(), json!([{"role": "user", "content": text}]));
283    }
284
285    // Add session ID
286    state.insert("session_id".to_string(), json!(ctx.session_id()));
287
288    state
289}
290
291/// Default output mapper - creates events from state
292fn default_output_mapper(state: &State) -> Vec<Event> {
293    let mut events = Vec::new();
294
295    // Try to get output from common fields
296    let output_text = state
297        .get("output")
298        .and_then(|v| v.as_str())
299        .or_else(|| state.get("result").and_then(|v| v.as_str()))
300        .or_else(|| {
301            state
302                .get("messages")
303                .and_then(|v| v.as_array())
304                .and_then(|arr| arr.last())
305                .and_then(|msg| msg.get("content"))
306                .and_then(|c| c.as_str())
307        });
308
309    let text = if let Some(text) = output_text {
310        text.to_string()
311    } else {
312        // Return the full state as JSON
313        serde_json::to_string_pretty(state).unwrap_or_default()
314    };
315
316    let mut event = Event::new("graph_output");
317    event.set_content(Content::new("assistant").with_text(&text));
318    events.push(event);
319
320    events
321}
322
323/// Builder for GraphAgent
324pub struct GraphAgentBuilder {
325    name: String,
326    description: String,
327    schema: StateSchema,
328    nodes: Vec<Arc<dyn Node>>,
329    edges: Vec<Edge>,
330    checkpointer: Option<Arc<dyn Checkpointer>>,
331    interrupt_before: Vec<String>,
332    interrupt_after: Vec<String>,
333    recursion_limit: usize,
334    max_concurrency: Option<usize>,
335    input_mapper: Option<InputMapper>,
336    output_mapper: Option<OutputMapper>,
337    before_callback: Option<BeforeAgentCallback>,
338    after_callback: Option<AfterAgentCallback>,
339    timeout_policies: HashMap<String, TimeoutPolicy>,
340    default_timeout: Option<TimeoutPolicy>,
341    deferred_configs: HashMap<String, DeferredNodeConfig>,
342    #[cfg(feature = "node-cache")]
343    cache_policies: HashMap<String, crate::cache::NodeCachePolicy>,
344}
345
346impl GraphAgentBuilder {
347    /// Create a new builder
348    pub fn new(name: &str) -> Self {
349        Self {
350            name: name.to_string(),
351            description: String::new(),
352            schema: StateSchema::simple(&["input", "output", "messages"]),
353            nodes: vec![],
354            edges: vec![],
355            checkpointer: None,
356            interrupt_before: vec![],
357            interrupt_after: vec![],
358            recursion_limit: 100,
359            max_concurrency: None,
360            input_mapper: None,
361            output_mapper: None,
362            before_callback: None,
363            after_callback: None,
364            timeout_policies: HashMap::new(),
365            default_timeout: None,
366            deferred_configs: HashMap::new(),
367            #[cfg(feature = "node-cache")]
368            cache_policies: HashMap::new(),
369        }
370    }
371
372    /// Set description
373    pub fn description(mut self, desc: &str) -> Self {
374        self.description = desc.to_string();
375        self
376    }
377
378    /// Set state schema
379    pub fn state_schema(mut self, schema: StateSchema) -> Self {
380        self.schema = schema;
381        self
382    }
383
384    /// Add channels to state schema
385    pub fn channels(mut self, channels: &[&str]) -> Self {
386        self.schema = StateSchema::simple(channels);
387        self
388    }
389
390    /// Add a node
391    pub fn node<N: Node + 'static>(mut self, node: N) -> Self {
392        self.nodes.push(Arc::new(node));
393        self
394    }
395
396    /// Add a function as a node
397    pub fn node_fn<F, Fut>(mut self, name: &str, func: F) -> Self
398    where
399        F: Fn(NodeContext) -> Fut + Send + Sync + 'static,
400        Fut: Future<Output = Result<NodeOutput>> + Send + 'static,
401    {
402        self.nodes.push(Arc::new(FunctionNode::new(name, func)));
403        self
404    }
405
406    /// Add a direct edge
407    pub fn edge(mut self, source: &str, target: &str) -> Self {
408        let target =
409            if target == END { EdgeTarget::End } else { EdgeTarget::Node(target.to_string()) };
410
411        if source == START {
412            let entry_idx = self.edges.iter().position(|e| matches!(e, Edge::Entry { .. }));
413            match entry_idx {
414                Some(idx) => {
415                    if let Edge::Entry { targets } = &mut self.edges[idx]
416                        && let EdgeTarget::Node(node) = &target
417                        && !targets.contains(node)
418                    {
419                        targets.push(node.clone());
420                    }
421                }
422                None => {
423                    if let EdgeTarget::Node(node) = target {
424                        self.edges.push(Edge::Entry { targets: vec![node] });
425                    }
426                }
427            }
428        } else {
429            self.edges.push(Edge::Direct { source: source.to_string(), target });
430        }
431
432        self
433    }
434
435    /// Add a conditional edge
436    pub fn conditional_edge<F, I>(mut self, source: &str, router: F, targets: I) -> Self
437    where
438        F: Fn(&State) -> String + Send + Sync + 'static,
439        I: IntoIterator<Item = (&'static str, &'static str)>,
440    {
441        let targets_map: HashMap<String, EdgeTarget> = targets
442            .into_iter()
443            .map(|(k, v)| {
444                let target =
445                    if v == END { EdgeTarget::End } else { EdgeTarget::Node(v.to_string()) };
446                (k.to_string(), target)
447            })
448            .collect();
449
450        self.edges.push(Edge::Conditional {
451            source: source.to_string(),
452            router: Arc::new(router),
453            targets: targets_map,
454        });
455
456        self
457    }
458
459    /// Set checkpointer
460    pub fn checkpointer<C: Checkpointer + 'static>(mut self, checkpointer: C) -> Self {
461        self.checkpointer = Some(Arc::new(checkpointer));
462        self
463    }
464
465    /// Set checkpointer with Arc
466    pub fn checkpointer_arc(mut self, checkpointer: Arc<dyn Checkpointer>) -> Self {
467        self.checkpointer = Some(checkpointer);
468        self
469    }
470
471    /// Set nodes to interrupt before
472    pub fn interrupt_before(mut self, nodes: &[&str]) -> Self {
473        self.interrupt_before = nodes.iter().map(|s| s.to_string()).collect();
474        self
475    }
476
477    /// Set nodes to interrupt after
478    pub fn interrupt_after(mut self, nodes: &[&str]) -> Self {
479        self.interrupt_after = nodes.iter().map(|s| s.to_string()).collect();
480        self
481    }
482
483    /// Set recursion limit
484    /// Cap how many nodes execute concurrently within one super-step.
485    ///
486    /// See [`CompiledGraph::with_max_concurrency`](crate::graph::CompiledGraph::with_max_concurrency).
487    pub fn max_concurrency(mut self, limit: usize) -> Self {
488        self.max_concurrency = Some(limit.max(1));
489        self
490    }
491
492    pub fn recursion_limit(mut self, limit: usize) -> Self {
493        self.recursion_limit = limit;
494        self
495    }
496
497    /// Set a timeout policy for a specific node.
498    ///
499    /// The policy is applied when the named node executes, enforcing
500    /// wall-clock and/or idle timeouts with the configured recovery action.
501    ///
502    /// # Example
503    ///
504    /// ```rust,ignore
505    /// use std::time::Duration;
506    /// use adk_graph::timeout::{TimeoutPolicy, OnTimeout};
507    ///
508    /// let agent = GraphAgent::builder("my_graph")
509    ///     .node_timeout("slow_node", TimeoutPolicy {
510    ///         run_timeout: Some(Duration::from_secs(10)),
511    ///         idle_timeout: None,
512    ///         on_timeout: OnTimeout::Fail,
513    ///     })
514    ///     .build()?;
515    /// ```
516    pub fn node_timeout(mut self, node_name: &str, policy: TimeoutPolicy) -> Self {
517        self.timeout_policies.insert(node_name.to_string(), policy);
518        self
519    }
520
521    /// Set a default timeout policy applied to all nodes without an explicit override.
522    ///
523    /// Nodes that have a per-node policy set via [`node_timeout`](Self::node_timeout)
524    /// will use their specific policy instead of this default.
525    ///
526    /// # Example
527    ///
528    /// ```rust,ignore
529    /// use std::time::Duration;
530    /// use adk_graph::timeout::{TimeoutPolicy, OnTimeout};
531    ///
532    /// let agent = GraphAgent::builder("my_graph")
533    ///     .default_timeout(TimeoutPolicy {
534    ///         run_timeout: Some(Duration::from_secs(30)),
535    ///         idle_timeout: Some(Duration::from_secs(5)),
536    ///         on_timeout: OnTimeout::Skip,
537    ///     })
538    ///     .build()?;
539    /// ```
540    pub fn default_timeout(mut self, policy: TimeoutPolicy) -> Self {
541        self.default_timeout = Some(policy);
542        self
543    }
544
545    /// Add a deferred (fan-in barrier) node to the graph.
546    ///
547    /// A deferred node waits for all upstream parallel paths to complete before
548    /// executing. The provided function is wrapped as a [`FunctionNode`] and the
549    /// [`DeferredNodeConfig`] controls how upstream outputs are merged and how
550    /// long the node waits for all paths.
551    ///
552    /// # Arguments
553    ///
554    /// * `name` - The name of the deferred node.
555    /// * `func` - The async function to execute once all upstream paths complete.
556    /// * `config` - Configuration controlling merge strategy and fan-in timeout.
557    ///
558    /// # Example
559    ///
560    /// ```rust,ignore
561    /// use std::time::Duration;
562    /// use adk_graph::deferred::{DeferredNodeConfig, MergeStrategy};
563    /// use adk_graph::node::NodeOutput;
564    ///
565    /// let agent = GraphAgent::builder("scatter_gather")
566    ///     .deferred_node("aggregator", |_ctx| async {
567    ///         Ok(NodeOutput::new().with_update("status", serde_json::json!("merged")))
568    ///     }, DeferredNodeConfig {
569    ///         merge_strategy: MergeStrategy::Collect,
570    ///         fan_in_timeout: Some(Duration::from_secs(30)),
571    ///     })
572    ///     .build()?;
573    /// ```
574    /// Configure fan-in for a node already added with [`node`](Self::node).
575    ///
576    /// [`deferred_node`](Self::deferred_node) both adds and configures a node, so
577    /// a custom `Node` added through `node` had no way to set a merge strategy or
578    /// a fan-in timeout.
579    ///
580    /// A node reached by more than one unconditional edge is deferred
581    /// automatically; this overrides that default.
582    pub fn mark_deferred(mut self, name: &str, config: DeferredNodeConfig) -> Self {
583        self.deferred_configs.insert(name.to_string(), config);
584        self
585    }
586
587    pub fn deferred_node<F, Fut>(mut self, name: &str, func: F, config: DeferredNodeConfig) -> Self
588    where
589        F: Fn(NodeContext) -> Fut + Send + Sync + 'static,
590        Fut: Future<Output = Result<NodeOutput>> + Send + 'static,
591    {
592        self.nodes.push(Arc::new(FunctionNode::new(name, func)));
593        self.deferred_configs.insert(name.to_string(), config);
594        self
595    }
596
597    /// Set a cache policy for a specific node.
598    ///
599    /// When a node has a cache policy, its execution results are cached keyed
600    /// by a blake3 hash of the node name and input state. Subsequent executions
601    /// with identical inputs return the cached result without re-executing the
602    /// node.
603    ///
604    /// # Arguments
605    ///
606    /// * `name` — the name of the node to cache
607    /// * `policy` — the cache policy specifying backend and TTL
608    ///
609    /// # Example
610    ///
611    /// ```rust,ignore
612    /// use std::time::Duration;
613    /// use adk_graph::cache::{CacheBackend, NodeCachePolicy};
614    ///
615    /// let agent = GraphAgent::builder("cached_graph")
616    ///     .node_cache("expensive_node", NodeCachePolicy {
617    ///         backend: CacheBackend::InMemory { max_entries: 128 },
618    ///         ttl: Some(Duration::from_secs(300)),
619    ///     })
620    ///     .build()?;
621    /// ```
622    #[cfg(feature = "node-cache")]
623    pub fn node_cache(mut self, name: &str, policy: crate::cache::NodeCachePolicy) -> Self {
624        self.cache_policies.insert(name.to_string(), policy);
625        self
626    }
627
628    /// Set custom input mapper
629    pub fn input_mapper<F>(mut self, mapper: F) -> Self
630    where
631        F: Fn(&dyn InvocationContext) -> State + Send + Sync + 'static,
632    {
633        self.input_mapper = Some(Arc::new(mapper));
634        self
635    }
636
637    /// Set custom output mapper
638    pub fn output_mapper<F>(mut self, mapper: F) -> Self
639    where
640        F: Fn(&State) -> Vec<Event> + Send + Sync + 'static,
641    {
642        self.output_mapper = Some(Arc::new(mapper));
643        self
644    }
645
646    /// Set before agent callback
647    pub fn before_agent_callback<F, Fut>(mut self, callback: F) -> Self
648    where
649        F: Fn(Arc<dyn InvocationContext>) -> Fut + Send + Sync + 'static,
650        Fut: Future<Output = adk_core::Result<()>> + Send + 'static,
651    {
652        self.before_callback = Some(Arc::new(move |ctx| Box::pin(callback(ctx))));
653        self
654    }
655
656    /// Set after agent callback
657    ///
658    /// Note: The callback receives a cloned Event to avoid lifetime issues.
659    pub fn after_agent_callback<F, Fut>(mut self, callback: F) -> Self
660    where
661        F: Fn(Arc<dyn InvocationContext>, Event) -> Fut + Send + Sync + 'static,
662        Fut: Future<Output = adk_core::Result<()>> + Send + 'static,
663    {
664        self.after_callback = Some(Arc::new(move |ctx, event| {
665            let event_clone = event.clone();
666            Box::pin(callback(ctx, event_clone))
667        }));
668        self
669    }
670
671    /// Add an action node to the graph.
672    ///
673    /// Wraps the `ActionNodeConfig` in an `ActionNodeExecutor` and registers it
674    /// as a node. If the config is a `SwitchNodeConfig`, conditional edges are
675    /// also auto-registered from the switch conditions.
676    #[cfg(feature = "action")]
677    pub fn action_node(mut self, config: adk_action::ActionNodeConfig) -> Self {
678        use crate::action::ActionNodeExecutor;
679
680        // If this is a Switch node, register conditional edges
681        if let adk_action::ActionNodeConfig::Switch(ref switch_config) = config {
682            let conditions = switch_config.conditions.clone();
683            let eval_mode = switch_config.evaluation_mode.clone();
684            let default_branch = switch_config.default_branch.clone();
685            let source = config.standard().id.clone();
686
687            let mut targets_map: HashMap<String, EdgeTarget> = HashMap::new();
688            for condition in &conditions {
689                targets_map.insert(
690                    condition.output_port.clone(),
691                    EdgeTarget::Node(condition.output_port.clone()),
692                );
693            }
694            if let Some(ref default) = default_branch {
695                let target = if default == END {
696                    EdgeTarget::End
697                } else {
698                    EdgeTarget::Node(default.clone())
699                };
700                targets_map.insert(default.clone(), target);
701            }
702            targets_map.insert(END.to_string(), EdgeTarget::End);
703
704            let router = Arc::new(move |state: &State| -> String {
705                match crate::action::switch::evaluate_switch_conditions(
706                    &conditions,
707                    state,
708                    &eval_mode,
709                    default_branch.as_deref(),
710                ) {
711                    Ok(ports) => ports.into_iter().next().unwrap_or_else(|| END.to_string()),
712                    Err(_) => END.to_string(),
713                }
714            });
715
716            self.edges.push(Edge::Conditional { source, router, targets: targets_map });
717        }
718
719        let executor = ActionNodeExecutor::new(config);
720        self.nodes.push(Arc::new(executor));
721        self
722    }
723
724    /// Build the GraphAgent
725    pub fn build(self) -> Result<GraphAgent> {
726        // Build the graph
727        let mut graph = StateGraph::new(self.schema);
728
729        // Add nodes
730        for node in self.nodes {
731            graph.nodes.insert(node.name().to_string(), node);
732        }
733
734        // Add edges
735        graph.edges = self.edges;
736
737        // Compile
738        let mut compiled = graph.compile()?;
739
740        // Configure
741        if let Some(cp) = self.checkpointer {
742            compiled.checkpointer = Some(cp);
743        }
744        compiled.interrupt_before = self.interrupt_before.into_iter().collect();
745        compiled.interrupt_after = self.interrupt_after.into_iter().collect();
746        compiled.recursion_limit = self.recursion_limit;
747        compiled.max_concurrency = self.max_concurrency;
748        compiled.timeout_policies = self.timeout_policies;
749        compiled.default_timeout = self.default_timeout;
750        compiled.deferred_configs = self.deferred_configs;
751
752        #[cfg(feature = "node-cache")]
753        {
754            compiled.cache_policies = self.cache_policies;
755        }
756
757        Ok(GraphAgent {
758            name: self.name,
759            description: self.description,
760            graph: Arc::new(compiled),
761            input_mapper: self.input_mapper.unwrap_or(Arc::new(default_input_mapper)),
762            output_mapper: self.output_mapper.unwrap_or(Arc::new(default_output_mapper)),
763            before_callback: self.before_callback,
764            after_callback: self.after_callback,
765        })
766    }
767}
768
769#[cfg(test)]
770mod tests {
771    use super::*;
772    use serde_json::json;
773
774    #[tokio::test]
775    async fn test_graph_agent_builder() {
776        let agent = GraphAgent::builder("test")
777            .description("Test agent")
778            .channels(&["value"])
779            .node_fn("set", |_ctx| async { Ok(NodeOutput::new().with_update("value", json!(42))) })
780            .edge(START, "set")
781            .edge("set", END)
782            .build()
783            .unwrap();
784
785        assert_eq!(agent.name(), "test");
786        assert_eq!(agent.description(), "Test agent");
787
788        // Test direct invocation
789        let result = agent.invoke(State::new(), ExecutionConfig::new("test")).await.unwrap();
790
791        assert_eq!(result.get("value"), Some(&json!(42)));
792
793        let topology = agent.topology().expect("graph topology");
794        assert_eq!(topology.root, "test");
795        assert_eq!(topology.coordinator, "test");
796        assert_eq!(topology.members.len(), 2);
797        assert_eq!(topology.relationships.len(), 1);
798        assert_eq!(topology.relationships[0].from, "test");
799        assert_eq!(topology.relationships[0].to, "set");
800        assert_eq!(topology.relationships[0].kind, AgentRelationshipKind::Flow);
801    }
802}
adk_graph/agent.rs

adk_graph/
agent.rs
