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