Skip to main content

zen_engine/
engine.rs

1use crate::decision::Decision;
2use crate::decision_graph::graph::{DecisionGraphResponse, EvaluationTrace};
3use crate::error::ContentKindError;
4use crate::loader::{ClosureLoader, DynamicLoader, LoaderResponse, LoaderResult, NoopLoader};
5use crate::model::{DecisionContent, GraphContent};
6use crate::nodes::custom::{DynamicCustomNode, NoopCustomNode};
7use crate::nodes::function::http_handler::DynamicHttpHandler;
8use crate::policy::runtime::{CompiledEntry, CompiledSet};
9use crate::{CompileFailure, EvaluationError};
10use arc_swap::ArcSwapOption;
11use serde_json::Value;
12use std::fmt::Debug;
13use std::future::Future;
14use std::sync::Arc;
15use strum::{EnumString, IntoStaticStr};
16use zen_expression::variable::Variable;
17
18/// Structure used for generating and evaluating JDM decisions
19#[derive(Clone)]
20pub struct DecisionEngine {
21    loader: DynamicLoader,
22    adapter: DynamicCustomNode,
23    http_handler: DynamicHttpHandler,
24    compiled: Arc<ArcSwapOption<CompiledSet>>,
25}
26
27impl Debug for DecisionEngine {
28    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29        f.debug_struct("DecisionEngine")
30            .field("loader", &self.loader)
31            .field("adapter", &self.adapter)
32            .field("http_handler", &self.http_handler)
33            .finish()
34    }
35}
36
37#[derive(Debug, Clone, Copy)]
38pub struct EvaluationOptions {
39    pub trace: bool,
40    pub max_depth: u8,
41}
42
43impl Default for EvaluationOptions {
44    fn default() -> Self {
45        Self {
46            trace: false,
47            max_depth: 10,
48        }
49    }
50}
51
52#[derive(Debug, Clone, Copy)]
53pub struct EvaluationSerializedOptions {
54    pub trace: EvaluationTraceKind,
55    pub max_depth: u8,
56}
57
58impl Default for EvaluationSerializedOptions {
59    fn default() -> Self {
60        Self {
61            trace: EvaluationTraceKind::None,
62            max_depth: 10,
63        }
64    }
65}
66
67#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, EnumString, IntoStaticStr)]
68#[strum(serialize_all = "camelCase")]
69pub enum EvaluationTraceKind {
70    #[default]
71    None,
72    Default,
73    String,
74    Reference,
75    ReferenceString,
76}
77
78impl EvaluationTraceKind {
79    pub fn serialize_trace(&self, trace: &Variable) -> Value {
80        match self {
81            EvaluationTraceKind::None => Value::Null,
82            EvaluationTraceKind::Default => serde_json::to_value(&trace).unwrap_or_default(),
83            EvaluationTraceKind::String => {
84                Value::String(serde_json::to_string(&trace).unwrap_or_default())
85            }
86            EvaluationTraceKind::Reference => {
87                serde_json::to_value(&trace.serialize_ref()).unwrap_or_default()
88            }
89            EvaluationTraceKind::ReferenceString => {
90                Value::String(serde_json::to_string(&trace.serialize_ref()).unwrap_or_default())
91            }
92        }
93    }
94}
95
96impl Default for DecisionEngine {
97    fn default() -> Self {
98        Self {
99            loader: Arc::new(NoopLoader::default()),
100            adapter: Arc::new(NoopCustomNode::default()),
101            http_handler: None,
102            compiled: Arc::new(ArcSwapOption::empty()),
103        }
104    }
105}
106
107impl DecisionEngine {
108    pub fn new(loader: DynamicLoader, adapter: DynamicCustomNode) -> Self {
109        Self {
110            loader,
111            adapter,
112            http_handler: None,
113            compiled: Arc::new(ArcSwapOption::empty()),
114        }
115    }
116
117    pub fn with_adapter(mut self, adapter: DynamicCustomNode) -> Self {
118        self.adapter = adapter;
119        self.compiled = Arc::new(ArcSwapOption::empty());
120        self
121    }
122
123    pub fn with_loader(mut self, loader: DynamicLoader) -> Self {
124        self.loader = loader;
125        self.compiled = Arc::new(ArcSwapOption::empty());
126        self
127    }
128
129    pub fn with_http_handler(mut self, http_handler: DynamicHttpHandler) -> Self {
130        self.http_handler = http_handler;
131        self.compiled = Arc::new(ArcSwapOption::empty());
132        self
133    }
134
135    pub fn with_closure_loader<F, O>(mut self, loader: F) -> Self
136    where
137        F: Fn(String) -> O + Sync + Send + 'static,
138        O: Future<Output = LoaderResponse> + Send,
139    {
140        self.loader = Arc::new(ClosureLoader::new(loader));
141        self.compiled = Arc::new(ArcSwapOption::empty());
142        self
143    }
144
145    pub fn compile(&self) -> Vec<CompileFailure> {
146        let Some(keys) = self.loader.keys() else {
147            return Vec::new();
148        };
149
150        let set = CompiledSet::build_sync(&self.loader, &keys);
151
152        let failures = set.failures().to_vec();
153        self.compiled.store(Some(Arc::new(set)));
154        failures
155    }
156
157    pub fn compile_failures(&self) -> Vec<CompileFailure> {
158        self.compiled
159            .load_full()
160            .map(|set| set.failures().to_vec())
161            .unwrap_or_default()
162    }
163
164    /// Evaluates a decision through loader using a key
165    pub async fn evaluate<K>(
166        &self,
167        key: K,
168        context: Variable,
169    ) -> Result<DecisionGraphResponse, Box<EvaluationError>>
170    where
171        K: AsRef<str>,
172    {
173        self.evaluate_with_opts(key, context, Default::default())
174            .await
175    }
176
177    /// Evaluates a decision through loader using a key with advanced options
178    pub async fn evaluate_with_opts<K>(
179        &self,
180        key: K,
181        context: Variable,
182        options: EvaluationOptions,
183    ) -> Result<DecisionGraphResponse, Box<EvaluationError>>
184    where
185        K: AsRef<str>,
186    {
187        let key_str = key.as_ref();
188        if let Some(set) = self.compiled.load_full() {
189            if let Some(entry) = set.get(key_str) {
190                return match entry {
191                    CompiledEntry::Policy(artifact) => artifact
192                        .evaluate_entry(key_str, context, options.trace)
193                        .map(|r| DecisionGraphResponse {
194                            performance: format!("{:.1?}", r.duration),
195                            result: r.output,
196                            trace: r.trace.map(EvaluationTrace::Policy),
197                        })
198                        .map_err(|e| Box::new(EvaluationError::Policy(e))),
199                    CompiledEntry::Graph(graph) => {
200                        self.decision_from_graph(graph)
201                            .evaluate_with_opts(context, options)
202                            .await
203                    }
204                };
205            }
206        }
207        let content = self.loader.load(key_str).await?;
208        match content.as_ref() {
209            DecisionContent::Graph(_) => {
210                let decision = self.decision_from_graph_arc(content);
211                decision.evaluate_with_opts(context, options).await
212            }
213            DecisionContent::Policy(_) => {
214                crate::policy::runtime::evaluate_policy(
215                    &self.loader,
216                    key_str,
217                    content,
218                    context,
219                    options,
220                )
221                .await
222            }
223        }
224    }
225
226    pub async fn evaluate_serialized<K>(
227        &self,
228        key: K,
229        context: Variable,
230        options: EvaluationSerializedOptions,
231    ) -> Result<Value, Value>
232    where
233        K: AsRef<str>,
234    {
235        let key_str = key.as_ref();
236        if let Some(set) = self.compiled.load_full() {
237            if let Some(entry) = set.get(key_str) {
238                match entry {
239                    CompiledEntry::Policy(artifact) => {
240                        let trace_mode = options.trace;
241                        let trace = options.trace != EvaluationTraceKind::None;
242                        return match artifact.evaluate_entry(key_str, context, trace) {
243                            Ok(r) => {
244                                let response = DecisionGraphResponse {
245                                    performance: format!("{:.1?}", r.duration),
246                                    result: r.output,
247                                    trace: r.trace.map(EvaluationTrace::Policy),
248                                };
249                                Ok(response
250                                    .serialize_with_mode(serde_json::value::Serializer, trace_mode)
251                                    .unwrap_or_default())
252                            }
253                            Err(e) => {
254                                let err = EvaluationError::Policy(e);
255                                Err(err
256                                    .serialize_with_mode(serde_json::value::Serializer, trace_mode)
257                                    .unwrap_or_default())
258                            }
259                        };
260                    }
261                    CompiledEntry::Graph(graph) => {
262                        return self
263                            .decision_from_graph(graph)
264                            .evaluate_serialized(context, options)
265                            .await;
266                    }
267                }
268            }
269        }
270        let content = self
271            .loader
272            .load(key_str)
273            .await
274            .map_err(|err| Value::String(err.to_string()))?;
275
276        match content.as_ref() {
277            DecisionContent::Graph(_) => {
278                let decision = self.decision_from_graph_arc(content);
279                decision.evaluate_serialized(context, options).await
280            }
281            DecisionContent::Policy(_) => {
282                let inner_opts = EvaluationOptions {
283                    trace: options.trace != EvaluationTraceKind::None,
284                    max_depth: options.max_depth,
285                };
286                let trace_mode = options.trace;
287                let response = crate::policy::runtime::evaluate_policy(
288                    &self.loader,
289                    key_str,
290                    content,
291                    context,
292                    inner_opts,
293                )
294                .await;
295                match response {
296                    Ok(ok) => Ok(ok
297                        .serialize_with_mode(serde_json::value::Serializer, trace_mode)
298                        .unwrap_or_default()),
299                    Err(err) => Err(err
300                        .serialize_with_mode(serde_json::value::Serializer, trace_mode)
301                        .unwrap_or_default()),
302                }
303            }
304        }
305    }
306
307    fn decision_from_graph_arc(&self, content: Arc<DecisionContent>) -> Decision {
308        let graph: Arc<GraphContent> = match Arc::try_unwrap(content) {
309            Ok(DecisionContent::Graph(g)) => g,
310            Err(arc) => match arc.as_ref() {
311                DecisionContent::Graph(g) => g.clone(),
312                DecisionContent::Policy(_) => {
313                    panic!("decision_from_graph_arc called with Policy variant")
314                }
315            },
316            Ok(DecisionContent::Policy(_)) => {
317                panic!("decision_from_graph_arc called with Policy variant")
318            }
319        };
320        self.decision_from_graph(graph)
321    }
322
323    fn decision_from_graph(&self, graph: Arc<GraphContent>) -> Decision {
324        Decision::from(graph)
325            .with_loader(self.loader.clone())
326            .with_adapter(self.adapter.clone())
327            .with_http_handler(self.http_handler.clone())
328    }
329
330    /// Creates a decision from DecisionContent, exists for easier binding creation
331    pub fn create_decision(
332        &self,
333        content: Arc<DecisionContent>,
334    ) -> Result<Decision, ContentKindError> {
335        match content.as_ref() {
336            DecisionContent::Graph(_) => Ok(self.decision_from_graph_arc(content)),
337            DecisionContent::Policy(_) => Err(ContentKindError {
338                expected: "graph",
339                got: "policy",
340            }),
341        }
342    }
343
344    /// Retrieves a decision based on the loader
345    pub async fn get_decision(
346        &self,
347        key: &str,
348    ) -> LoaderResult<Result<Decision, ContentKindError>> {
349        let content = self.loader.load(key).await?;
350        Ok(self.create_decision(content))
351    }
352    pub fn loader(&self) -> DynamicLoader {
353        self.loader.clone()
354    }
355}