zen_engine/nodes/decision/
mod.rs1use crate::decision_graph::graph::{DecisionGraph, DecisionGraphConfig};
2use crate::nodes::{NodeContext, NodeContextExt, NodeError, NodeHandler, NodeResult};
3use crate::EvaluationError;
4use std::cell::RefCell;
5use std::ops::Deref;
6use std::rc::Rc;
7use zen_types::decision::{DecisionNodeContent, TransformAttributes};
8use zen_types::variable::{ToVariable, Variable};
9
10#[derive(Debug, Clone, Default)]
11pub struct DecisionNodeHandler {
12 decision_graph: Rc<RefCell<Option<DecisionGraph>>>,
13}
14
15pub type DecisionNodeData = DecisionNodeContent;
16pub type DecisionNodeTrace = Variable;
17
18impl NodeHandler for DecisionNodeHandler {
19 type NodeData = DecisionNodeData;
20 type TraceData = DecisionNodeTrace;
21
22 fn transform_attributes(
23 &self,
24 ctx: &NodeContext<Self::NodeData, Self::TraceData>,
25 ) -> Option<TransformAttributes> {
26 Some(ctx.node.transform_attributes.clone())
27 }
28
29 async fn after_transform_attributes(
30 &self,
31 _ctx: &NodeContext<Self::NodeData, Self::TraceData>,
32 ) -> Result<(), NodeError> {
33 if let Some(graph) = self.decision_graph.borrow_mut().as_mut() {
34 graph.reset_graph();
35 };
36
37 Ok(())
38 }
39
40 async fn handle(&self, ctx: NodeContext<Self::NodeData, Self::TraceData>) -> NodeResult {
41 let loader = ctx.extensions.loader();
42 let sub_decision = loader.load(ctx.node.key.deref()).await.node_context(&ctx)?;
43 let sub_kind = sub_decision.kind();
44 let Some(sub_graph) = sub_decision.into_graph_arc() else {
45 return ctx.error(format!(
46 "sub-decision '{}' is a {sub_kind}, expected graph",
47 ctx.node.key
48 ));
49 };
50
51 let mut decision_graph_ref = self.decision_graph.borrow_mut();
52 let decision_graph = match decision_graph_ref.as_mut() {
53 Some(dg) => dg,
54 None => {
55 let dg = DecisionGraph::try_new(DecisionGraphConfig {
56 content: sub_graph,
57 extensions: ctx.extensions.clone(),
58 trace: ctx.config.trace,
59 iteration: ctx.iteration + 1,
60 max_depth: ctx.config.max_depth,
61 })
62 .node_context(&ctx)?;
63
64 *decision_graph_ref = Some(dg);
65 match decision_graph_ref.as_mut() {
66 Some(dg) => dg,
67 None => return ctx.error("Failed to initialize decision graph".to_string()),
68 }
69 }
70 };
71
72 let evaluate_result = Box::pin(decision_graph.evaluate(ctx.input.clone())).await;
73 match evaluate_result {
74 Ok(result) => {
75 ctx.trace(|trace| {
76 *trace = result
77 .trace
78 .and_then(|t| t.into_graph())
79 .as_ref()
80 .map(|m| m.to_variable())
81 .unwrap_or(Variable::Null);
82 });
83
84 ctx.success(result.result)
85 }
86 Err(err) => {
87 if let EvaluationError::NodeError { trace, .. } = err.deref() {
88 ctx.trace(|t| *t = trace.to_variable());
89 }
90
91 ctx.error(err.to_string())
92 }
93 }
94 }
95}