1use std::collections::HashMap;
2
3use agentrs_core::{AgentError, AgentOutput, Result};
4
5#[derive(Clone, Default)]
7pub struct AgentGraph {
8 nodes: HashMap<String, String>,
9 edges: Vec<GraphEdge>,
10 entry: Option<String>,
11}
12
13#[derive(Clone, Default)]
15pub struct AgentGraphBuilder {
16 graph: AgentGraph,
17}
18
19#[derive(Clone)]
21pub struct GraphEdge {
22 pub from: String,
24 pub to: String,
26 pub condition: EdgeCondition,
28}
29
30#[derive(Clone)]
32pub enum EdgeCondition {
33 Always,
35 Contains(String),
37 End,
39}
40
41impl AgentGraph {
42 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 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 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 pub fn entry(mut self, entry: impl Into<String>) -> Self {
93 self.graph.entry = Some(entry.into());
94 self
95 }
96
97 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}