ri_agent_graph/
router.rs

1use crate::config::GraphConfig;
2use crate::error::Result;
3use crate::state::AgentState;
4use async_trait::async_trait;
5use std::future::Future;
6use std::pin::Pin;
7
8/// Router output determines where execution goes next.
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub enum RouterOutput {
11    /// Route to a single node, or end execution (None)
12    Next(Option<String>),
13    /// Fan-out to multiple nodes simultaneously
14    FanOut(Vec<String>),
15}
16
17impl From<Option<String>> for RouterOutput {
18    fn from(opt: Option<String>) -> Self {
19        RouterOutput::Next(opt)
20    }
21}
22
23impl From<String> for RouterOutput {
24    fn from(s: String) -> Self {
25        RouterOutput::Next(Some(s))
26    }
27}
28
29impl From<Vec<String>> for RouterOutput {
30    fn from(v: Vec<String>) -> Self {
31        RouterOutput::FanOut(v)
32    }
33}
34
35/// Determines which node to visit next based on current state.
36#[async_trait]
37pub trait RoutingFunction: Send + Sync {
38    /// Returns routing decision.
39    /// Return `RouterOutput::Next(None)` to end execution.
40    async fn route(&self, state: &AgentState, config: &GraphConfig) -> Result<RouterOutput>;
41
42    /// Stable caller-visible identity for routing semantics. Implementations
43    /// should override this when the route behavior depends on configuration.
44    fn semantic_digest(&self) -> String {
45        std::any::type_name::<Self>().to_string()
46    }
47}
48
49/// Helper to create a router from an async function
50pub struct FnRouter<F>
51where
52    F: Fn(&AgentState, &GraphConfig) -> Pin<Box<dyn Future<Output = Result<RouterOutput>> + Send>>
53        + Send
54        + Sync,
55{
56    func: F,
57}
58
59impl<F> FnRouter<F>
60where
61    F: Fn(&AgentState, &GraphConfig) -> Pin<Box<dyn Future<Output = Result<RouterOutput>> + Send>>
62        + Send
63        + Sync,
64{
65    pub fn new(func: F) -> Self {
66        Self { func }
67    }
68}
69
70#[async_trait]
71impl<F> RoutingFunction for FnRouter<F>
72where
73    F: Fn(&AgentState, &GraphConfig) -> Pin<Box<dyn Future<Output = Result<RouterOutput>> + Send>>
74        + Send
75        + Sync,
76{
77    async fn route(&self, state: &AgentState, config: &GraphConfig) -> Result<RouterOutput> {
78        (self.func)(state, config).await
79    }
80}
81
82/// Helper macro to create a router from an async closure.
83///
84/// # Forms
85///
86/// ```ignore
87/// // Basic form (backward compatible) - body returns Result<Option<String>>
88/// router!(|state| async move {
89///     let value: i32 = state.get("value").await?;
90///     Ok(if value > 5 { Some("high".to_string()) } else { None })
91/// })
92///
93/// // With config - body returns Result<impl Into<RouterOutput>>
94/// router!(|state, config| async move {
95///     let value: i32 = state.get("value").await?;
96///     Ok(RouterOutput::FanOut(vec!["a".to_string(), "b".to_string()]))
97/// })
98/// ```
99#[macro_export]
100macro_rules! router {
101    // Form 1: |state| - backward compatible, returns Result<impl Into<RouterOutput>>
102    (|$state:ident| async move $body:block) => {
103        Box::new($crate::router::FnRouter::new(
104            |__state: &$crate::state::AgentState, __config: &$crate::config::GraphConfig| {
105                let $state = __state.clone();
106                let _ = __config;
107                Box::pin(async move {
108                    let __result = (|| async move { $body })().await;
109                    __result.map(::std::convert::Into::into)
110                })
111            },
112        ))
113    };
114    // Form 2: |state, config| - has access to config
115    (|$state:ident, $config:ident| async move $body:block) => {
116        Box::new($crate::router::FnRouter::new(
117            |__state: &$crate::state::AgentState, __config: &$crate::config::GraphConfig| {
118                let $state = __state.clone();
119                let $config = __config.clone();
120                Box::pin(async move {
121                    let __result = (|| async move { $body })().await;
122                    __result.map(::std::convert::Into::into)
123                })
124            },
125        ))
126    };
127}
ri_agent_graph/router.rs

ri_agent_graph/
router.rs
