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 mut decision_graph_ref = self.decision_graph.borrow_mut();
42
43        if decision_graph_ref.is_none() {
44            let loader = ctx.extensions.loader();
45            let sub_decision = loader.load(ctx.node.key.deref()).await.node_context(&ctx)?;
46            let sub_kind = sub_decision.kind();
47            let Some(sub_graph) = sub_decision.into_graph_arc() else {
48                return ctx.error(format!(
49                    "sub-decision '{}' is a {sub_kind}, expected graph",
50                    ctx.node.key
51                ));
52            };
53
54            let sub_graph =
55                if sub_graph.compiled_cache.is_some() && sub_graph.resolved_schemas.is_some() {
56                    sub_graph
57                } else {
58                    let mut owned = (*sub_graph).clone();
59                    owned.compile();
60                    let _ = owned.resolve_schemas(loader).await;
61                    std::sync::Arc::new(owned)
62                };
63
64            let mut extensions = ctx.extensions.clone();
65            extensions.compiled_cache = sub_graph.compiled_cache.clone();
66            extensions.dt_indexes = sub_graph.dt_indexes.clone();
67            extensions.validator_cache =
68                std::sync::Arc::new(std::cell::OnceCell::from(sub_graph.validator_cache.clone()));
69
70            let dg = DecisionGraph::try_new(DecisionGraphConfig {
71                content: sub_graph,
72                extensions,
73                trace: ctx.config.trace,
74                iteration: ctx.iteration + 1,
75                max_depth: ctx.config.max_depth,
76            })
77            .node_context(&ctx)?;
78
79            *decision_graph_ref = Some(dg);
80        }
81
82        let Some(decision_graph) = decision_graph_ref.as_mut() else {
83            return ctx.error("Failed to initialize decision graph".to_string());
84        };
85
86        decision_graph.set_parent_nodes(ctx.nodes.clone());
87
88        let evaluate_result = Box::pin(decision_graph.evaluate(ctx.input.clone())).await;
89        match evaluate_result {
90            Ok(result) => {
91                ctx.trace(|trace| {
92                    *trace = result
93                        .trace
94                        .and_then(|t| t.into_graph())
95                        .as_ref()
96                        .map(|m| m.to_variable())
97                        .unwrap_or(Variable::Null);
98                });
99
100                ctx.success(result.result)
101            }
102            Err(err) => {
103                if let EvaluationError::NodeError { trace, .. } = err.deref() {
104                    ctx.trace(|t| *t = trace.to_variable());
105                }
106
107                ctx.error(err.to_string())
108            }
109        }
110    }
111}
zen_engine/nodes/decision/mod.rs

zen_engine/nodes/decision/
mod.rs
