zen_engine/decision_graph/
graph.rs

1use crate::decision_graph::cleaner::VariableCleaner;
2use crate::decision_graph::schema_dict;
3use crate::decision_graph::tracer::NodeTracer;
4use crate::decision_graph::walker::{GraphWalker, NodeData, StableDiDecisionGraph};
5use crate::engine::EvaluationTraceKind;
6use crate::model::{DecisionNodeKind, GraphContent};
7use crate::nodes::custom::CustomNodeHandler;
8use crate::nodes::decision::DecisionNodeHandler;
9use crate::nodes::decision_table::DecisionTableNodeHandler;
10use crate::nodes::expression::ExpressionNodeHandler;
11use crate::nodes::function::FunctionNodeHandler;
12use crate::nodes::input::InputNodeHandler;
13use crate::nodes::output::OutputNodeHandler;
14use crate::nodes::transform_attributes::TransformAttributesExecution;
15use crate::nodes::{
16    NodeContext, NodeContextBase, NodeContextConfig, NodeDataType, NodeHandler,
17    NodeHandlerExtensions, NodeResponse, NodeResult, TraceDataType,
18};
19use crate::{DecisionGraphTrace, DecisionGraphValidationError, EvaluationError};
20use ahash::{HashMap, HashMapExt};
21use petgraph::algo::is_cyclic_directed;
22use petgraph::matrix_graph::Zero;
23use serde::ser::SerializeMap;
24use serde::{Serialize, Serializer};
25use std::cell::RefCell;
26use std::ops::Deref;
27use std::rc::Rc;
28use std::sync::Arc;
29use std::time::Instant;
30use zen_expression::variable::{ToVariable, Variable};
31use zen_types::decision::{DecisionNode, InputNodeContent, OutputNodeContent};
32
33#[derive(Debug)]
34pub struct DecisionGraph {
35    initial_graph: StableDiDecisionGraph,
36    graph: StableDiDecisionGraph,
37    config: DecisionGraphConfig,
38}
39
40#[derive(Debug)]
41pub struct DecisionGraphConfig {
42    pub content: Arc<GraphContent>,
43    pub trace: bool,
44    pub iteration: u8,
45    pub max_depth: u8,
46    pub extensions: NodeHandlerExtensions,
47}
48
49impl DecisionGraph {
50    pub fn try_new(config: DecisionGraphConfig) -> Result<Self, DecisionGraphValidationError> {
51        let graph = Self::build_graph(config.content.deref())?;
52        Ok(Self {
53            initial_graph: graph.clone(),
54            graph,
55            config,
56        })
57    }
58
59    fn build_graph(
60        content: &GraphContent,
61    ) -> Result<StableDiDecisionGraph, DecisionGraphValidationError> {
62        let mut graph = StableDiDecisionGraph::new();
63        let mut index_map = HashMap::with_capacity(content.nodes.len());
64
65        for node in &content.nodes {
66            let node_id = node.id.clone();
67            let node_index = graph.add_node(node.clone());
68
69            index_map.insert(node_id, node_index);
70        }
71
72        for edge in &content.edges {
73            let source_index = index_map.get(&edge.source_id).ok_or_else(|| {
74                DecisionGraphValidationError::MissingNode(edge.source_id.to_string())
75            })?;
76
77            let target_index = index_map.get(&edge.target_id).ok_or_else(|| {
78                DecisionGraphValidationError::MissingNode(edge.target_id.to_string())
79            })?;
80
81            graph.add_edge(*source_index, *target_index, edge.clone());
82        }
83
84        Ok(graph)
85    }
86
87    pub(crate) fn reset_graph(&mut self) {
88        self.graph = self.initial_graph.clone();
89    }
90
91    pub fn validate(&self) -> Result<(), DecisionGraphValidationError> {
92        let input_count = self
93            .graph
94            .node_weights()
95            .filter(|w| matches!(w.kind, DecisionNodeKind::InputNode { .. }))
96            .count();
97        if input_count != 1 {
98            return Err(DecisionGraphValidationError::InvalidInputCount(
99                input_count as u32,
100            ));
101        }
102
103        if is_cyclic_directed(&self.graph) {
104            return Err(DecisionGraphValidationError::CyclicGraph);
105        }
106
107        Ok(())
108    }
109
110    async fn validation_schema(
111        &self,
112        schema: Option<&serde_json::Value>,
113    ) -> Result<Option<(serde_json::Value, u64)>, String> {
114        let Some(schema) = schema else {
115            return Ok(None);
116        };
117        if !schema_dict::schema_references_dictionary(schema) {
118            return Ok(None);
119        }
120        let dictionaries = schema_dict::load_import_dictionaries(
121            self.config.extensions.loader(),
122            &self.config.content.imports,
123        )
124        .await?;
125        schema_dict::resolve_schema(schema, &dictionaries).map(Some)
126    }
127
128    fn build_node_context(&self, node: &DecisionNode, input: Variable) -> NodeContextBase {
129        NodeContextBase {
130            id: node.id.clone(),
131            name: node.name.clone(),
132            input,
133            extensions: self.config.extensions.clone(),
134            iteration: self.config.iteration,
135            trace: match self.config.trace {
136                true => Some(RefCell::new(Variable::Null)),
137                false => None,
138            },
139            config: NodeContextConfig {
140                max_depth: self.config.max_depth,
141                trace: self.config.trace,
142                ..Default::default()
143            },
144        }
145    }
146
147    pub async fn evaluate(
148        &mut self,
149        context: Variable,
150    ) -> Result<DecisionGraphResponse, Box<EvaluationError>> {
151        let root_start = Instant::now();
152
153        self.validate()?;
154        if self.config.iteration >= self.config.max_depth {
155            return Err(Box::new(EvaluationError::DepthLimitExceeded));
156        }
157
158        let mut walker = GraphWalker::new(&self.graph);
159        let mut tracer = NodeTracer::new(self.config.trace);
160
161        while let Some(nid) = walker.next(&mut self.graph, tracer.trace_callback()) {
162            if let Some(_) = walker.get_node_data(nid) {
163                continue;
164            }
165
166            let node = &self.graph[nid];
167            let start = Instant::now();
168            let (input, input_trace) = walker.incoming_node_data(&self.graph, nid, true);
169            let mut base_ctx = self.build_node_context(node.deref(), input);
170
171            let node_execution = match &node.kind {
172                DecisionNodeKind::InputNode { content } => {
173                    base_ctx.input = context.clone();
174                    match self.validation_schema(content.schema.as_deref()).await {
175                        Err(message) => base_ctx.error(message),
176                        Ok(None) => handle_node(base_ctx, content.clone(), InputNodeHandler).await,
177                        Ok(Some((schema, salt))) => {
178                            base_ctx.config.validation_salt = salt;
179                            let resolved = InputNodeContent {
180                                schema: Some(Arc::new(schema)),
181                            };
182                            handle_node(base_ctx, resolved, InputNodeHandler).await
183                        }
184                    }
185                }
186                DecisionNodeKind::OutputNode { content } => {
187                    match self.validation_schema(content.schema.as_deref()).await {
188                        Err(message) => base_ctx.error(message),
189                        Ok(None) => handle_node(base_ctx, content.clone(), OutputNodeHandler).await,
190                        Ok(Some((schema, salt))) => {
191                            base_ctx.config.validation_salt = salt;
192                            let resolved = OutputNodeContent {
193                                schema: Some(Arc::new(schema)),
194                            };
195                            handle_node(base_ctx, resolved, OutputNodeHandler).await
196                        }
197                    }
198                }
199                DecisionNodeKind::SwitchNode { .. } => Ok(NodeResponse {
200                    output: input_trace.clone(),
201                    trace_data: None,
202                }),
203                DecisionNodeKind::FunctionNode { content } => {
204                    handle_node(base_ctx, content.clone(), FunctionNodeHandler).await
205                }
206                DecisionNodeKind::DecisionNode { content } => {
207                    handle_node(base_ctx, content.clone(), DecisionNodeHandler::default()).await
208                }
209                DecisionNodeKind::DecisionTableNode { content } => {
210                    handle_node(base_ctx, content.clone(), DecisionTableNodeHandler).await
211                }
212                DecisionNodeKind::ExpressionNode { content } => {
213                    handle_node(base_ctx, content.clone(), ExpressionNodeHandler).await
214                }
215                DecisionNodeKind::CustomNode { content } => {
216                    handle_node(base_ctx, content.clone(), CustomNodeHandler).await
217                }
218            };
219
220            tracer.record_execution(node.deref(), input_trace, &node_execution, start.elapsed());
221
222            let output = match node_execution {
223                Ok(ok) => ok.output,
224                Err(err) => {
225                    let mut cleaner = VariableCleaner::new();
226                    let trace = tracer.into_traces();
227                    if let Some(t) = &trace {
228                        t.values().for_each(|v| {
229                            cleaner.clean(&v.input);
230                            cleaner.clean(&v.output);
231                            if let Some(td) = &v.trace_data {
232                                cleaner.clean(td);
233                            }
234                        })
235                    }
236
237                    return Err(Box::new(EvaluationError::NodeError {
238                        node_id: err.node_id,
239                        source: err.source,
240                        trace: trace.map(|t| t.to_variable()),
241                    }));
242                }
243            };
244
245            walker.set_node_data(
246                nid,
247                NodeData {
248                    name: Rc::from(node.name.deref()),
249                    data: output,
250                },
251            );
252
253            // Terminate once Output node is reached
254            if matches!(node.kind, DecisionNodeKind::OutputNode { .. }) {
255                break;
256            }
257        }
258
259        let result = walker.ending_variables(&self.graph);
260        let trace = tracer.into_traces();
261
262        if self.config.iteration.is_zero() {
263            let mut cleaner = VariableCleaner::new();
264            cleaner.clean(&result);
265            if let Some(t) = &trace {
266                t.values().for_each(|v| {
267                    cleaner.clean(&v.input);
268                    cleaner.clean(&v.output);
269                    if let Some(td) = &v.trace_data {
270                        cleaner.clean(td);
271                    }
272                })
273            }
274        }
275
276        Ok(DecisionGraphResponse {
277            performance: format!("{:.1?}", root_start.elapsed()),
278            result,
279            trace: trace.map(EvaluationTrace::Graph),
280        })
281    }
282}
283
284#[derive(Debug, Clone, Serialize)]
285#[serde(untagged)]
286pub enum EvaluationTrace {
287    Graph(HashMap<Arc<str>, DecisionGraphTrace>),
288    Policy(crate::policy::Trace),
289}
290
291impl EvaluationTrace {
292    pub fn as_graph(&self) -> Option<&HashMap<Arc<str>, DecisionGraphTrace>> {
293        match self {
294            Self::Graph(m) => Some(m),
295            Self::Policy(_) => None,
296        }
297    }
298
299    pub fn into_graph(self) -> Option<HashMap<Arc<str>, DecisionGraphTrace>> {
300        match self {
301            Self::Graph(m) => Some(m),
302            Self::Policy(_) => None,
303        }
304    }
305}
306
307#[derive(Debug, Clone, Serialize)]
308#[serde(rename_all = "camelCase")]
309pub struct DecisionGraphResponse {
310    pub performance: String,
311    pub result: Variable,
312    #[serde(skip_serializing_if = "Option::is_none")]
313    pub trace: Option<EvaluationTrace>,
314}
315
316impl DecisionGraphResponse {
317    pub fn serialize_with_mode<S>(
318        &self,
319        serializer: S,
320        mode: EvaluationTraceKind,
321    ) -> Result<S::Ok, S::Error>
322    where
323        S: Serializer,
324    {
325        let mut map = serializer.serialize_map(None)?;
326        map.serialize_entry("performance", &self.performance)?;
327        map.serialize_entry("result", &self.result)?;
328        if let Some(trace) = &self.trace {
329            match trace {
330                EvaluationTrace::Graph(graph_trace) => {
331                    map.serialize_entry(
332                        "trace",
333                        &mode.serialize_trace(&graph_trace.to_variable()),
334                    )?;
335                }
336                EvaluationTrace::Policy(policy_trace) => match mode {
337                    EvaluationTraceKind::String | EvaluationTraceKind::ReferenceString => {
338                        map.serialize_entry(
339                            "trace",
340                            &serde_json::to_string(policy_trace).unwrap_or_default(),
341                        )?;
342                    }
343                    _ => {
344                        map.serialize_entry("trace", policy_trace)?;
345                    }
346                },
347            }
348        }
349
350        map.end()
351    }
352}
353
354async fn handle_node<NodeData, TraceData, NodeHandlerType>(
355    base_ctx: NodeContextBase,
356    content: NodeData,
357    handler: NodeHandlerType,
358) -> NodeResult
359where
360    TraceData: TraceDataType,
361    NodeData: NodeDataType,
362    NodeHandlerType: NodeHandler<NodeData = NodeData, TraceData = TraceData>,
363{
364    let ctx = NodeContext::<NodeData, TraceData>::from_base(base_ctx.clone(), content);
365    if let Some(transform_attributes) = handler.transform_attributes(&ctx) {
366        return transform_attributes
367            .run_with(base_ctx, move |input, has_more| {
368                let handler = handler.clone();
369                let mut new_ctx = ctx.clone();
370                new_ctx.input = input;
371
372                async move {
373                    match has_more {
374                        false => handler.handle(new_ctx).await,
375                        true => {
376                            let result = handler.handle(new_ctx.clone()).await;
377                            handler.after_transform_attributes(&new_ctx).await?;
378                            result
379                        }
380                    }
381                }
382            })
383            .await;
384    }
385
386    handler.handle(ctx).await
387}
zen_engine/decision_graph/graph.rs

zen_engine/decision_graph/
graph.rs
