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