1use atman_dsl::ast::{Arg, Expr, FanoutCollect, FlowDecl, Ident, Node, Stmt};
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
5pub struct FlowGraph {
6 pub flow_name: String,
7 pub root: Vec<StaticNode>,
8}
9
10#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
11pub struct StaticNode {
12 pub node_id: String,
13 pub kind: NodeKind,
14 pub label: String,
15 pub children: Vec<StaticNode>,
16}
17
18#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
19#[serde(tag = "type", rename_all = "snake_case")]
20pub enum NodeKind {
21 Llm { model: Option<String> },
22 ToolCall { path: String },
23 Fanout { collect: FanoutMode },
24 UserConfirm,
25 Subflow { name: String },
26 Message { role: String },
27 FixUntilTest,
28 When { condition_preview: String },
29 Return,
30}
31
32#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
33#[serde(rename_all = "snake_case")]
34pub enum FanoutMode {
35 All,
36 First,
37}
38
39impl From<FanoutCollect> for FanoutMode {
40 fn from(c: FanoutCollect) -> Self {
41 match c {
42 FanoutCollect::All => FanoutMode::All,
43 FanoutCollect::First => FanoutMode::First,
44 }
45 }
46}
47
48pub fn extract_graph(flow: &FlowDecl) -> FlowGraph {
49 let mut root = Vec::new();
50 for (i, stmt) in flow.body.iter().enumerate() {
51 extract_stmt(stmt, &format!("{i}"), &mut root);
52 }
53 FlowGraph {
54 flow_name: flow.name.name.clone(),
55 root,
56 }
57}
58
59fn extract_stmt(stmt: &Stmt, prefix: &str, out: &mut Vec<StaticNode>) {
60 match stmt {
61 Stmt::Bind { value, .. } => extract_expr(value, prefix, out),
62 Stmt::Expr(expr) => extract_expr(expr, prefix, out),
63 Stmt::Return { value } => {
64 extract_expr(value, &format!("{prefix}.v"), out);
65 out.push(StaticNode {
66 node_id: prefix.to_string(),
67 kind: NodeKind::Return,
68 label: "return".into(),
69 children: Vec::new(),
70 });
71 }
72 Stmt::When { cond, body } => {
73 let mut inner = Vec::new();
74 for (i, s) in body.iter().enumerate() {
75 extract_stmt(s, &format!("{prefix}.{i}"), &mut inner);
76 }
77 out.push(StaticNode {
78 node_id: prefix.to_string(),
79 kind: NodeKind::When {
80 condition_preview: format_expr_short(cond),
81 },
82 label: format!("when {}", format_expr_short(cond)),
83 children: inner,
84 });
85 }
86 Stmt::Watch(_) => {}
87 }
88}
89
90fn extract_expr(expr: &Expr, prefix: &str, out: &mut Vec<StaticNode>) {
91 match expr {
92 Expr::Node(node) => extract_node(node, prefix, out),
93 Expr::Pipe { lhs, rhs } => {
94 extract_expr(lhs, &format!("{prefix}.l"), out);
95 extract_expr(rhs, &format!("{prefix}.r"), out);
96 }
97 _ => {}
98 }
99}
100
101fn extract_node(node: &Node, prefix: &str, out: &mut Vec<StaticNode>) {
102 let (kind, label, children) = match node {
103 Node::Llm { kwargs } => {
104 let model = kwargs
105 .iter()
106 .find(|(k, _)| k.name == "model")
107 .and_then(|(_, v)| match v {
108 Expr::Literal(atman_dsl::ast::Literal::Str(s)) => Some(s.clone()),
109 _ => None,
110 });
111 let label = match &model {
112 Some(m) => format!("llm({m})"),
113 None => "llm".to_string(),
114 };
115 (NodeKind::Llm { model }, label, Vec::new())
116 }
117 Node::ToolCall { path, .. } => {
118 let path_str = path
119 .iter()
120 .map(|p| p.name.clone())
121 .collect::<Vec<_>>()
122 .join(".");
123 let label = format!("⟶ {path_str}");
124 (NodeKind::ToolCall { path: path_str }, label, Vec::new())
125 }
126 Node::Fanout { items, collect } => {
127 let mut branch_children = Vec::new();
128 for (i, item) in items.iter().enumerate() {
129 extract_expr(item, &format!("{prefix}.branch[{i}]"), &mut branch_children);
130 }
131 let label = format!("fanout ×{}", items.len());
132 (
133 NodeKind::Fanout {
134 collect: (*collect).into(),
135 },
136 label,
137 branch_children,
138 )
139 }
140 Node::UserConfirm { .. } => (NodeKind::UserConfirm, "user_confirm".into(), Vec::new()),
141 Node::Subflow { name, args } => {
142 let label = format!(
143 "subflow({}{})",
144 name.name,
145 if args.is_empty() { "" } else { ", …" }
146 );
147 (
148 NodeKind::Subflow {
149 name: name.name.clone(),
150 },
151 label,
152 Vec::new(),
153 )
154 }
155 Node::Message { role, args } => {
156 let role_str = match role {
157 atman_dsl::ast::MessageRole::User => "user",
158 atman_dsl::ast::MessageRole::Assistant => "assistant",
159 atman_dsl::ast::MessageRole::System => "system",
160 atman_dsl::ast::MessageRole::Tool => "tool",
161 };
162 let _ = args;
163 (
164 NodeKind::Message {
165 role: role_str.into(),
166 },
167 format!("{role_str}_msg"),
168 Vec::new(),
169 )
170 }
171 Node::FixUntilTestPasses { .. } => {
172 (NodeKind::FixUntilTest, "fix_until_test".into(), Vec::new())
173 }
174 };
175 out.push(StaticNode {
176 node_id: prefix.to_string(),
177 kind,
178 label,
179 children,
180 });
181}
182
183fn format_expr_short(expr: &Expr) -> String {
184 match expr {
185 Expr::Literal(atman_dsl::ast::Literal::Bool(b)) => b.to_string(),
186 Expr::Literal(atman_dsl::ast::Literal::Str(s)) => format!("\"{s}\""),
187 Expr::Literal(atman_dsl::ast::Literal::Int(i)) => i.to_string(),
188 Expr::Literal(atman_dsl::ast::Literal::Float(f)) => f.to_string(),
189 Expr::Ident(id) => id.name.clone(),
190 Expr::Binary { op, left, right } => {
191 let sym = match op {
192 atman_dsl::ast::BinOp::Eq => "==",
193 atman_dsl::ast::BinOp::Ne => "!=",
194 atman_dsl::ast::BinOp::Lt => "<",
195 atman_dsl::ast::BinOp::Le => "<=",
196 atman_dsl::ast::BinOp::Gt => ">",
197 atman_dsl::ast::BinOp::Ge => ">=",
198 atman_dsl::ast::BinOp::And => "and",
199 atman_dsl::ast::BinOp::Or => "or",
200 atman_dsl::ast::BinOp::Add => "+",
201 atman_dsl::ast::BinOp::Sub => "-",
202 atman_dsl::ast::BinOp::Mul => "*",
203 atman_dsl::ast::BinOp::Div => "/",
204 atman_dsl::ast::BinOp::Mod => "%",
205 };
206 format!(
207 "{} {} {}",
208 format_expr_short(left),
209 sym,
210 format_expr_short(right)
211 )
212 }
213 _ => "…".into(),
214 }
215}
216
217#[allow(dead_code)]
218fn _unused_ident(_: &Ident, _: &[Arg]) {}
219
220#[cfg(test)]
221mod tests {
222 use super::*;
223 use atman_dsl::parse::parse_file;
224
225 fn parse_first_flow(src: &str) -> FlowDecl {
226 let file = parse_file(src).expect("parse ok");
227 file.flows.into_iter().next().expect("has flow")
228 }
229
230 #[test]
231 fn extracts_llm_only_flow() {
232 let src = r#"flow smoke() -> string {
233 x = llm {
234 model: "glm"
235 messages: []
236 }
237 return x
238 }"#;
239 let flow = parse_first_flow(src);
240 let g = extract_graph(&flow);
241 assert_eq!(g.flow_name, "smoke");
242 let kinds: Vec<_> = g.root.iter().map(|n| n.kind.clone()).collect();
243 assert!(matches!(kinds[0], NodeKind::Llm { .. }));
244 assert!(matches!(kinds.last(), Some(NodeKind::Return)));
245 }
246
247 #[test]
248 fn extracts_fanout_branches_from_example() {
249 let src = std::fs::read_to_string(concat!(
250 env!("CARGO_MANIFEST_DIR"),
251 "/../../examples/look_into.at"
252 ))
253 .expect("example loadable");
254 let file = parse_file(&src).expect("parse ok");
255 let flow = file
256 .flows
257 .iter()
258 .find(|f| f.name.name == "look_into")
259 .expect("has flow");
260 let g = extract_graph(flow);
261 let fanout = g
262 .root
263 .iter()
264 .find(|n| matches!(n.kind, NodeKind::Fanout { .. }))
265 .expect("has fanout");
266 assert!(fanout.children.len() >= 2);
267 }
268
269 #[test]
270 fn extracts_when_body() {
271 let src = r#"flow t() -> string {
272 when true {
273 a = llm { model: "m" messages: [] }
274 }
275 return "x"
276 }"#;
277 let flow = parse_first_flow(src);
278 let g = extract_graph(&flow);
279 let when = g
280 .root
281 .iter()
282 .find(|n| matches!(n.kind, NodeKind::When { .. }));
283 assert!(when.is_some());
284 assert_eq!(when.unwrap().children.len(), 1);
285 }
286
287 #[test]
288 fn simple_return_only_flow() {
289 let src = r#"flow t() -> string { return "hi" }"#;
290 let flow = parse_first_flow(src);
291 let g = extract_graph(&flow);
292 assert_eq!(g.root.len(), 1);
293 assert!(matches!(g.root[0].kind, NodeKind::Return));
294 }
295}