Skip to main content

zen_engine/nodes/decision/
mod.rs

1use 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}