use std::collections::HashMap; use std::rc::Rc; use std::sync::Arc; use std::time::Instant; use anyhow::anyhow; use petgraph::algo::is_cyclic_directed; use serde::ser::SerializeMap; use serde::{Deserialize, Serialize, Serializer}; use serde_json::Value; use thiserror::Error; use crate::handler::custom_node_adapter::{CustomNodeAdapter, CustomNodeRequest}; use crate::handler::decision::DecisionHandler; use crate::handler::expression::ExpressionHandler; use crate::handler::function::function::{Function, FunctionConfig}; use crate::handler::function::module::console::ConsoleListener; use crate::handler::function::module::zen::ZenListener; use crate::handler::function::FunctionHandler; use crate::handler::function_v1; use crate::handler::function_v1::runtime::create_runtime; use crate::handler::node::NodeRequest; use crate::handler::table::zen::DecisionTableHandler; use crate::handler::traversal::{GraphWalker, StableDiDecisionGraph}; use crate::loader::DecisionLoader; use crate::model::{DecisionContent, DecisionNodeKind, FunctionNodeContent}; use crate::{EvaluationError, NodeError}; pub struct DecisionGraph<'a, L: DecisionLoader + 'static, A: CustomNodeAdapter + 'static> { graph: StableDiDecisionGraph<'a>, adapter: Arc, loader: Arc, trace: bool, max_depth: u8, iteration: u8, runtime: Option>, } pub struct DecisionGraphConfig<'a, L: DecisionLoader + 'static, A: CustomNodeAdapter + 'static> { pub loader: Arc, pub adapter: Arc, pub content: &'a DecisionContent, pub trace: bool, pub iteration: u8, pub max_depth: u8, } impl<'a, L: DecisionLoader + 'static, A: CustomNodeAdapter + 'static> DecisionGraph<'a, L, A> { pub fn try_new( config: DecisionGraphConfig<'a, L, A>, ) -> Result { let content = config.content; let mut graph = StableDiDecisionGraph::new(); let mut index_map = HashMap::new(); for node in &content.nodes { let node_id = node.id.clone(); let node_index = graph.add_node(node); index_map.insert(node_id, node_index); } for (_, edge) in content.edges.iter().enumerate() { let source_index = index_map.get(&edge.source_id).ok_or_else(|| { DecisionGraphValidationError::MissingNode(edge.source_id.to_string()) })?; let target_index = index_map.get(&edge.target_id).ok_or_else(|| { DecisionGraphValidationError::MissingNode(edge.target_id.to_string()) })?; graph.add_edge(source_index.clone(), target_index.clone(), edge); } Ok(Self { graph, iteration: config.iteration, trace: config.trace, loader: config.loader, adapter: config.adapter, max_depth: config.max_depth, runtime: None, }) } pub(crate) fn with_function(mut self, runtime: Option>) -> Self { self.runtime = runtime; self } async fn get_or_insert_function(&mut self) -> anyhow::Result> { if let Some(function) = &self.runtime { return Ok(function.clone()); } let function = Function::create(FunctionConfig { listeners: Some(vec![ Box::new(ConsoleListener), Box::new(ZenListener { loader: self.loader.clone(), adapter: self.adapter.clone(), }), ]), }) .await .map_err(|err| anyhow!(err.to_string()))?; let rc_function = Rc::new(function); self.runtime.replace(rc_function.clone()); Ok(rc_function) } pub fn validate(&self) -> Result<(), DecisionGraphValidationError> { let input_count = self.node_kind_count(DecisionNodeKind::InputNode); if input_count != 1 { return Err(DecisionGraphValidationError::InvalidInputCount( input_count as u32, )); } let output_count = self.node_kind_count(DecisionNodeKind::OutputNode); if output_count < 1 { return Err(DecisionGraphValidationError::InvalidOutputCount( output_count as u32, )); } if is_cyclic_directed(&self.graph) { return Err(DecisionGraphValidationError::CyclicGraph); } Ok(()) } fn node_kind_count(&self, kind: DecisionNodeKind) -> usize { self.graph .node_weights() .filter(|weight| weight.kind == kind) .count() } pub async fn evaluate(&mut self, context: &Value) -> Result { let root_start = Instant::now(); self.validate().map_err(|e| NodeError { node_id: "".to_string(), source: anyhow!(e), })?; if self.iteration >= self.max_depth { return Err(NodeError { node_id: "".to_string(), source: anyhow!(EvaluationError::DepthLimitExceeded), }); } let mut walker = GraphWalker::new(&self.graph); let mut node_traces = self.trace.then(|| HashMap::default()); while let Some(nid) = walker.next( &mut self.graph, self.trace.then_some(|trace: DecisionGraphTrace| { if let Some(nt) = &mut node_traces { nt.insert(trace.id.clone(), trace); }; }), ) { if let Some(_) = walker.get_node_data(nid) { continue; } let node = self.graph[nid]; let start = Instant::now(); macro_rules! trace { ($data: tt) => { if let Some(nt) = &mut node_traces { nt.insert(node.id.clone(), DecisionGraphTrace $data); }; }; } match &node.kind { DecisionNodeKind::InputNode => { walker.set_node_data(nid, context.clone()); trace!({ input: Value::Null, output: Value::Null, name: node.name.clone(), id: node.id.clone(), performance: None, trace_data: None, }); } DecisionNodeKind::OutputNode => { trace!({ input: Value::Null, output: Value::Null, name: node.name.clone(), id: node.id.clone(), performance: None, trace_data: None, }); return Ok(DecisionGraphResponse { result: walker.incoming_node_data(&self.graph, nid, false), performance: format!("{:?}", root_start.elapsed()), trace: node_traces, }); } DecisionNodeKind::SwitchNode { .. } => { let input_data = walker.incoming_node_data(&self.graph, nid, false); walker.set_node_data(nid, input_data); } DecisionNodeKind::FunctionNode { content } => { let function = self.get_or_insert_function().await.map_err(|e| NodeError { source: e.into(), node_id: node.id.clone(), })?; let mut node_request = NodeRequest { node, iteration: self.iteration, input: walker.incoming_node_data(&self.graph, nid, true), }; let mut res = match content { FunctionNodeContent::Version2(_) => FunctionHandler::new( function, self.trace, self.iteration, self.max_depth, ) .handle(&node_request) .await .map_err(|e| NodeError { source: e.into(), node_id: node.id.clone(), })?, FunctionNodeContent::Version1(_) => { let runtime = create_runtime().map_err(|e| NodeError { source: e.into(), node_id: node.id.clone(), })?; function_v1::FunctionHandler::new(self.trace, runtime) .handle(&node_request) .await .map_err(|e| NodeError { source: e.into(), node_id: node.id.clone(), })? } }; trim_nodes(&mut node_request.input); trim_nodes(&mut res.output); trace!({ input: node_request.input, output: res.output.clone(), name: node.name.clone(), id: node.id.clone(), performance: Some(format!("{:?}", start.elapsed())), trace_data: res.trace_data, }); walker.set_node_data(nid, res.output); } DecisionNodeKind::DecisionNode { .. } => { let mut node_request = NodeRequest { node, iteration: self.iteration, input: walker.incoming_node_data(&self.graph, nid, true), }; let mut res = DecisionHandler::new( self.trace, self.max_depth, self.loader.clone(), self.adapter.clone(), self.runtime.clone(), ) .handle(&node_request) .await .map_err(|e| NodeError { source: e.into(), node_id: node.id.to_string(), })?; trim_nodes(&mut node_request.input); trim_nodes(&mut res.output); trace!({ input: node_request.input, output: res.output.clone(), name: node.name.clone(), id: node.id.clone(), performance: Some(format!("{:?}", start.elapsed())), trace_data: res.trace_data, }); walker.set_node_data(nid, res.output); } DecisionNodeKind::DecisionTableNode { .. } => { let mut node_request = NodeRequest { node, iteration: self.iteration, input: walker.incoming_node_data(&self.graph, nid, true), }; let mut res = DecisionTableHandler::new(self.trace) .handle(&node_request) .await .map_err(|e| NodeError { node_id: node.id.clone(), source: e.into(), })?; trim_nodes(&mut node_request.input); trim_nodes(&mut res.output); trace!({ input: node_request.input, output: res.output.clone(), name: node.name.clone(), id: node.id.clone(), performance: Some(format!("{:?}", start.elapsed())), trace_data: res.trace_data, }); walker.set_node_data(nid, res.output); } DecisionNodeKind::ExpressionNode { .. } => { let mut node_request = NodeRequest { node, iteration: self.iteration, input: walker.incoming_node_data(&self.graph, nid, true), }; let mut res = ExpressionHandler::new(self.trace) .handle(&node_request) .await .map_err(|e| NodeError { node_id: node.id.clone(), source: e.into(), })?; trim_nodes(&mut node_request.input); trim_nodes(&mut res.output); trace!({ input: node_request.input, output: res.output.clone(), name: node.name.clone(), id: node.id.clone(), performance: Some(format!("{:?}", start.elapsed())), trace_data: res.trace_data, }); walker.set_node_data(nid, res.output); } DecisionNodeKind::CustomNode { .. } => { let mut node_request = NodeRequest { node, iteration: self.iteration, input: walker.incoming_node_data(&self.graph, nid, true), }; let mut res = self .adapter .handle(CustomNodeRequest::try_from(&node_request).unwrap()) .await .map_err(|e| NodeError { node_id: node.id.clone(), source: e.into(), })?; trim_nodes(&mut node_request.input); trim_nodes(&mut res.output); trace!({ input: node_request.input, output: res.output.clone(), name: node.name.clone(), id: node.id.clone(), performance: Some(format!("{:?}", start.elapsed())), trace_data: res.trace_data, }); walker.set_node_data(nid, res.output); } } } Err(NodeError { node_id: "".to_string(), source: anyhow!("Graph did not halt. Missing output node."), }) } } fn trim_nodes(val: &mut Value) { if let Some(obj) = val.as_object_mut() { obj.remove("$nodes"); } } #[derive(Debug, Error)] pub enum DecisionGraphValidationError { #[error("Invalid input node count: {0}")] InvalidInputCount(u32), #[error("Invalid output node count: {0}")] InvalidOutputCount(u32), #[error("Cyclic graph detected")] CyclicGraph, #[error("Missing node")] MissingNode(String), } impl Serialize for DecisionGraphValidationError { fn serialize(&self, serializer: S) -> Result where S: Serializer, { let mut map = serializer.serialize_map(None)?; match &self { DecisionGraphValidationError::InvalidInputCount(count) => { map.serialize_entry("type", "invalidInputCount")?; map.serialize_entry("nodeCount", count)?; } DecisionGraphValidationError::InvalidOutputCount(count) => { map.serialize_entry("type", "invalidOutputCount")?; map.serialize_entry("nodeCount", count)?; } DecisionGraphValidationError::MissingNode(node_id) => { map.serialize_entry("type", "missingNode")?; map.serialize_entry("nodeId", node_id)?; } DecisionGraphValidationError::CyclicGraph => { map.serialize_entry("type", "cyclicGraph")?; } } map.end() } } #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct DecisionGraphResponse { pub performance: String, pub result: Value, #[serde(skip_serializing_if = "Option::is_none")] pub trace: Option>, } #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct DecisionGraphTrace { pub input: Value, pub output: Value, pub name: String, pub id: String, pub performance: Option, pub trace_data: Option, }