feat: custom node (#138)

* feat: custom node

* add zen template, expose $nodes and $root

* add custom handler to go and nodejs

* add support for python, improve bindings, add error to zen template

* fix: correct binding exports for nodejs and python

* fix benchmark

* improve rust api, trim template in zen templates

* update cargo action

* update expression version

* compile action

* fix action format

bindings/python/Cargo.toml

@@ -12,10 +12,12 @@ crate-type = ["cdylib"]
 [dependencies]
 async-trait = { workspace = true }
 anyhow = { workspace = true }
-pyo3 = { version = "0.20.2", features = ["anyhow", "serde"] }
-pythonize = "0.20.0"
+pyo3 = { version = "0.20", features = ["anyhow", "serde"] }
+pythonize = "0.20"
+json_dotpath = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 futures = { workspace = true }
 zen-engine = { path = "../../core/engine" }
-zen-expression = { path = "../../core/expression" }
+zen-expression = { path = "../../core/expression" }
+zen-template = { path = "../../core/template" }

bindings/python/index.py

@@ -6,6 +6,12 @@ def loader(key):
     with open("../../test-data/" + key, "r") as f:
         return f.read()
 
+def custom_handler(request):
+    p1 = request.get_field("prop1")
+    return {
+        "output": { "sum": p1 }
+    }
+
 # The test based on unittest module
 class ZenEngine(unittest.TestCase):
     def test_decision_using_loader(self):
@@ -41,6 +47,16 @@ class ZenEngine(unittest.TestCase):
         r = functionDecision.evaluate({"input": 15})
         self.assertEqual(r["result"]["output"], 30)
 
+    def test_engine_custom_handler(self):
+        engine = zen.ZenEngine({ "loader": loader, "customHandler": custom_handler })
+        r1 = engine.evaluate("custom.json", {"a": 10})
+        r2 = engine.evaluate("custom.json", {"a": 20})
+        r3 = engine.evaluate("custom.json", {"a": 30})
+
+        self.assertEqual(r1["result"]["sum"], 20)
+        self.assertEqual(r2["result"]["sum"], 30)
+        self.assertEqual(r3["result"]["sum"], 40)
+
     def test_evaluate_expression(self):
         result = zen.evaluate_expression("sum(a)", { "a": [1, 2, 3, 4] })
         self.assertEqual(result, 10)
@@ -49,5 +65,9 @@ class ZenEngine(unittest.TestCase):
         result = zen.evaluate_unary_expression("'FR', 'ES', 'GB'", { "$": "GB" })
         self.assertEqual(result, True)
 
+    def test_render_template(self):
+        result = zen.render_template("{{ a + b }}", { "a": 10, "b": 20 })
+        self.assertEqual(result, 30)
+
 # run the test
 unittest.main()

bindings/python/src/custom_node.rs

@@ -0,0 +1,42 @@
+use anyhow::anyhow;
+use pyo3::types::PyDict;
+use pyo3::{PyObject, Python};
+use pythonize::depythonize;
+
+use zen_engine::handler::custom_node_adapter::{CustomNodeAdapter, CustomNodeRequest};
+use zen_engine::handler::node::{NodeResponse, NodeResult};
+
+use crate::types::PyNodeRequest;
+
+#[derive(Default)]
+pub(crate) struct PyCustomNode(Option<PyObject>);
+
+impl From<PyObject> for PyCustomNode {
+    fn from(value: PyObject) -> Self {
+        Self(Some(value))
+    }
+}
+
+impl From<Option<PyObject>> for PyCustomNode {
+    fn from(value: Option<PyObject>) -> Self {
+        Self(value)
+    }
+}
+
+impl CustomNodeAdapter for PyCustomNode {
+    async fn handle(&self, request: CustomNodeRequest<'_>) -> NodeResult {
+        let Some(callable) = &self.0 else {
+            return Err(anyhow!("Custom node handler not provided"));
+        };
+
+        let content: NodeResponse = Python::with_gil(|py| {
+            let req = PyNodeRequest::from_request(py, request)?;
+            let result = callable.call1(py, (req,))?;
+
+            let dict = result.extract::<&PyDict>(py)?;
+            depythonize(dict)
+        })?;
+
+        Ok(content)
+    }
+}

bindings/python/src/decision.rs

@@ -1,19 +1,23 @@
-use crate::engine::PyZenEvaluateOptions;
-use crate::loader::PyDecisionLoader;
-use crate::value::PyValue;
+use std::sync::Arc;
+
 use anyhow::{anyhow, Context};
 use pyo3::types::PyDict;
 use pyo3::{pyclass, pymethods, PyObject, PyResult, Python, ToPyObject};
 use pythonize::depythonize;
-use std::sync::Arc;
+
 use zen_engine::{Decision, EvaluationOptions};
 
+use crate::custom_node::PyCustomNode;
+use crate::engine::PyZenEvaluateOptions;
+use crate::loader::PyDecisionLoader;
+use crate::value::PyValue;
+
 #[pyclass]
 #[pyo3(name = "ZenDecision")]
-pub struct PyZenDecision(pub(crate) Arc<Decision<PyDecisionLoader>>);
+pub struct PyZenDecision(pub(crate) Arc<Decision<PyDecisionLoader, PyCustomNode>>);
 
-impl From<Decision<PyDecisionLoader>> for PyZenDecision {
-    fn from(value: Decision<PyDecisionLoader>) -> Self {
+impl From<Decision<PyDecisionLoader, PyCustomNode>> for PyZenDecision {
+    fn from(value: Decision<PyDecisionLoader, PyCustomNode>) -> Self {
         Self(value.into())
     }
 }

bindings/python/src/engine.rs

@@ -1,3 +1,4 @@
+use crate::custom_node::PyCustomNode;
 use crate::decision::PyZenDecision;
 use crate::loader::PyDecisionLoader;
 use crate::value::PyValue;
@@ -13,7 +14,7 @@ use zen_engine::{DecisionEngine, EvaluationOptions};
 #[pyclass]
 #[pyo3(name = "ZenEngine")]
 pub struct PyZenEngine {
-    graph: Arc<DecisionEngine<PyDecisionLoader>>,
+    graph: Arc<DecisionEngine<PyDecisionLoader, PyCustomNode>>,
 }
 
 #[derive(Serialize, Deserialize)]
@@ -34,7 +35,11 @@ impl Default for PyZenEvaluateOptions {
 impl Default for PyZenEngine {
     fn default() -> Self {
         Self {
-            graph: DecisionEngine::new(PyDecisionLoader::default()).into(),
+            graph: DecisionEngine::new(
+                Arc::new(PyDecisionLoader::default()),
+                Arc::new(PyCustomNode::default()),
+            )
+            .into(),
         }
     }
 }
@@ -47,13 +52,22 @@ impl PyZenEngine {
             return Ok(Default::default());
         };
 
-        let Some(loader_any) = options.get_item("loader")? else {
-            return Ok(Default::default());
+        let loader = match options.get_item("loader")? {
+            Some(loader) => Some(Python::with_gil(|py| loader.to_object(py))),
+            None => None,
+        };
+
+        let custom_node = match options.get_item("customHandler")? {
+            Some(custom_node) => Some(Python::with_gil(|py| custom_node.to_object(py))),
+            None => None,
         };
 
-        let loader = Python::with_gil(|py| loader_any.to_object(py));
         Ok(Self {
-            graph: DecisionEngine::new(PyDecisionLoader::from(loader)).into(),
+            graph: DecisionEngine::new(
+                Arc::new(PyDecisionLoader::from(loader)),
+                Arc::new(PyCustomNode::from(custom_node)),
+            )
+            .into(),
         })
     }
 

bindings/python/src/expression.rs

@@ -33,3 +33,13 @@ pub fn evaluate_unary_expression(expression: String, ctx: &PyDict) -> PyResult<b
 
     Ok(result)
 }
+
+#[pyfunction]
+pub fn render_template(py: Python, template: String, ctx: &PyDict) -> PyResult<PyObject> {
+    let context: Value = depythonize(ctx).context("Failed to convert context")?;
+
+    let result = zen_template::render(template.as_str(), &context)
+        .map_err(|e| anyhow!(serde_json::to_string(&e).unwrap_or_else(|_| e.to_string())))?;
+
+    Ok(PyValue(result).to_object(py))
+}

bindings/python/src/lib.rs

@@ -1,13 +1,15 @@
 use crate::decision::PyZenDecision;
 use crate::engine::PyZenEngine;
-use crate::expression::{evaluate_expression, evaluate_unary_expression};
+use crate::expression::{evaluate_expression, evaluate_unary_expression, render_template};
 use pyo3::types::PyModule;
 use pyo3::{pymodule, wrap_pyfunction, PyResult, Python};
 
+mod custom_node;
 mod decision;
 mod engine;
 mod expression;
 mod loader;
+mod types;
 mod value;
 
 #[pymodule]
@@ -16,6 +18,7 @@ fn zen(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<PyZenDecision>()?;
     m.add_function(wrap_pyfunction!(evaluate_expression, m)?)?;
     m.add_function(wrap_pyfunction!(evaluate_unary_expression, m)?)?;
+    m.add_function(wrap_pyfunction!(render_template, m)?)?;
 
     Ok(())
 }

bindings/python/src/loader.rs

@@ -14,6 +14,12 @@ impl From<PyObject> for PyDecisionLoader {
     }
 }
 
+impl From<Option<PyObject>> for PyDecisionLoader {
+    fn from(value: Option<PyObject>) -> Self {
+        Self(value)
+    }
+}
+
 impl PyDecisionLoader {
     fn load_element(&self, key: &str) -> Result<Arc<DecisionContent>, anyhow::Error> {
         let Some(object) = &self.0 else {

bindings/python/src/types.rs

@@ -0,0 +1,118 @@
+use anyhow::{anyhow, Context};
+use json_dotpath::DotPaths;
+use pyo3::{pyclass, pymethods, PyObject, PyResult, Python, ToPyObject};
+use serde::Serialize;
+use serde_json::Value;
+
+use zen_engine::handler::custom_node_adapter::{
+    CustomDecisionNode as BaseCustomDecisionNode, CustomNodeRequest,
+};
+use zen_engine::handler::node::NodeResponse;
+
+use crate::value::{value_to_object, PyValue};
+
+#[derive(Serialize)]
+struct CustomDecisionNode {
+    pub id: String,
+    pub name: String,
+    pub kind: String,
+    pub config: Value,
+}
+
+impl From<BaseCustomDecisionNode<'_>> for CustomDecisionNode {
+    fn from(value: BaseCustomDecisionNode) -> Self {
+        Self {
+            id: value.id.to_string(),
+            name: value.name.to_string(),
+            kind: value.kind.to_string(),
+            config: value.config.clone(),
+        }
+    }
+}
+
+#[pyclass]
+pub struct PyNodeRequest {
+    inner_node: CustomDecisionNode,
+    inner_input: Value,
+
+    #[pyo3(get)]
+    pub input: PyObject,
+    #[pyo3(get)]
+    pub node: PyObject,
+}
+
+impl PyNodeRequest {
+    pub fn from_request(
+        py: Python,
+        value: CustomNodeRequest<'_>,
+    ) -> pythonize::Result<PyNodeRequest> {
+        let inner_node = value.node.into();
+        let node_val = serde_json::to_value(&inner_node).unwrap();
+
+        Ok(Self {
+            input: value_to_object(py, &value.input),
+            node: value_to_object(py, &node_val),
+
+            inner_input: value.input.clone(),
+            inner_node,
+        })
+    }
+}
+
+#[pymethods]
+impl PyNodeRequest {
+    fn get_field(&self, py: Python, path: String) -> PyResult<PyObject> {
+        let node_config = &self.inner_node.config;
+
+        let selected_value: Value = node_config
+            .dot_get(path.as_str())
+            .ok()
+            .flatten()
+            .context("Failed to find JSON path")?;
+        let Value::String(template) = selected_value else {
+            return Ok(PyValue(selected_value).to_object(py));
+        };
+
+        let template_value = zen_template::render(template.as_str(), &self.inner_input)
+            .map_err(|e| anyhow!(serde_json::to_string(&e).unwrap_or_else(|_| e.to_string())))?;
+
+        Ok(PyValue(template_value).to_object(py))
+    }
+
+    fn get_field_raw(&self, py: Python, path: String) -> PyResult<PyObject> {
+        let node_config = &self.inner_node.config;
+
+        let selected_value: Value = node_config
+            .dot_get(path.as_str())
+            .ok()
+            .flatten()
+            .context("Failed to find JSON path")?;
+
+        Ok(PyValue(selected_value).to_object(py))
+    }
+}
+
+#[pyclass]
+#[derive(Clone)]
+pub struct PyNodeResponse {
+    pub output: Value,
+    pub trace_data: Option<Value>,
+}
+
+impl From<NodeResponse> for PyNodeResponse {
+    fn from(value: NodeResponse) -> Self {
+        Self {
+            output: value.output,
+            trace_data: value.trace_data,
+        }
+    }
+}
+
+impl From<PyNodeResponse> for NodeResponse {
+    fn from(value: PyNodeResponse) -> Self {
+        Self {
+            output: value.output,
+            trace_data: value.trace_data,
+        }
+    }
+}

bindings/python/src/value.rs

@@ -6,7 +6,7 @@ use std::collections::HashMap;
 #[derive(Clone, Debug)]
 pub struct PyValue(pub Value);
 
-fn value_to_object(val: &Value, py: Python<'_>) -> PyObject {
+pub fn value_to_object(py: Python<'_>, val: &Value) -> PyObject {
     match val {
         Value::Null => py.None(),
         Value::Bool(b) => b.to_object(py),
@@ -18,11 +18,11 @@ fn value_to_object(val: &Value, py: Python<'_>) -> PyObject {
         }
         Value::String(s) => s.to_object(py),
         Value::Array(v) => {
-            let inner: Vec<_> = v.iter().map(|x| value_to_object(x, py)).collect();
+            let inner: Vec<_> = v.iter().map(|x| value_to_object(py, x)).collect();
             inner.to_object(py)
         }
         Value::Object(m) => {
-            let inner: HashMap<_, _> = m.iter().map(|(k, v)| (k, value_to_object(v, py))).collect();
+            let inner: HashMap<_, _> = m.iter().map(|(k, v)| (k, value_to_object(py, v))).collect();
             inner.to_object(py)
         }
     }
@@ -30,6 +30,6 @@ fn value_to_object(val: &Value, py: Python<'_>) -> PyObject {
 
 impl ToPyObject for PyValue {
     fn to_object(&self, py: Python<'_>) -> PyObject {
-        value_to_object(&self.0, py)
+        value_to_object(py, &self.0)
     }
 }