Skip to main content

agentrs_multi/
graph.rs

1use std::collections::HashMap;
2
3use agentrs_core::{AgentError, AgentOutput, Result};
4
5/// Directed agent graph.
6#[derive(Clone, Default)]
7pub struct AgentGraph {
8    nodes: HashMap<String, String>,
9    edges: Vec<GraphEdge>,
10    entry: Option<String>,
11}
12
13/// Builder for [`AgentGraph`].
14#[derive(Clone, Default)]
15pub struct AgentGraphBuilder {
16    graph: AgentGraph,
17}
18
19/// Graph edge between two agent nodes.
20#[derive(Clone)]
21pub struct GraphEdge {
22    /// Source node name.
23    pub from: String,
24    /// Target node name.
25    pub to: String,
26    /// Edge predicate.
27    pub condition: EdgeCondition,
28}
29
30/// Transition predicate for graph orchestration.
31#[derive(Clone)]
32pub enum EdgeCondition {
33    /// Always traverse this edge.
34    Always,
35    /// Traverse when output contains a keyword.
36    Contains(String),
37    /// Terminal edge marker.
38    End,
39}
40
41impl AgentGraph {
42    /// Starts building a graph.
43    pub fn builder() -> AgentGraphBuilder {
44        AgentGraphBuilder::default()
45    }
46
47    pub(crate) fn entry(&self) -> Result<&str> {
48        self.entry.as_deref().ok_or_else(|| {
49            AgentError::InvalidConfiguration("graph entry point not set".to_string())
50        })
51    }
52
53    pub(crate) fn next(&self, current: &str, output: &AgentOutput) -> Option<String> {
54        self.edges
55            .iter()
56            .find(|edge| {
57                edge.from == current
58                    && match &edge.condition {
59                        EdgeCondition::Always => true,
60                        EdgeCondition::Contains(keyword) => output.text.contains(keyword),
61                        EdgeCondition::End => false,
62                    }
63            })
64            .map(|edge| edge.to.clone())
65    }
66}
67
68impl AgentGraphBuilder {
69    /// Adds a node by name.
70    pub fn node(mut self, name: impl Into<String>) -> Self {
71        let name = name.into();
72        self.graph.nodes.insert(name.clone(), name);
73        self
74    }
75
76    /// Adds an edge.
77    pub fn edge(
78        mut self,
79        from: impl Into<String>,
80        to: impl Into<String>,
81        condition: EdgeCondition,
82    ) -> Self {
83        self.graph.edges.push(GraphEdge {
84            from: from.into(),
85            to: to.into(),
86            condition,
87        });
88        self
89    }
90
91    /// Sets the graph entry node.
92    pub fn entry(mut self, entry: impl Into<String>) -> Self {
93        self.graph.entry = Some(entry.into());
94        self
95    }
96
97    /// Finalizes the graph.
98    pub fn build(self) -> Result<AgentGraph> {
99        let entry = self.graph.entry()?;
100        if !self.graph.nodes.contains_key(entry) {
101            return Err(AgentError::InvalidConfiguration(format!(
102                "graph entry node '{entry}' is not registered"
103            )));
104        }
105        Ok(self.graph)
106    }
107}