feat: make in trace return The traceable execution object

tabkram · Dec 15, 2023 · c38efe9 · c38efe9
1 parent 726db14
commit c38efe9
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/src/engine/executionEngineDecorators.ts b/src/engine/executionEngineDecorators.ts
@@ -58,10 +58,11 @@ export function engine(options?: { id: string }): ClassDecorator {
 
 /**
  * A method decorator that enables tracing for the decorated method.
- * @param options - Trace options for the execution.
- * @returns A decorator function.
+ *
+ * @param {TraceOptions<Array<any>, O> | TraceOptions<Array<any>, O>['trace']} [options] - Optional tracing options.
+ * @returns {Function} - A decorator function.
  */
-export function run<O>(options?: TraceOptions<Array<any>, O>) {
+export function run<O>(options?: TraceOptions<Array<any>, O> | TraceOptions<Array<any>, O>['trace']) {
   return function (target: any, propertyKey: string, descriptor: PropertyDescriptor) {
     // Store the original method
     const originalMethod = descriptor.value;

diff --git a/src/trace/traceableExecution.ts b/src/trace/traceableExecution.ts
@@ -86,10 +86,17 @@ export class TraceableExecution {
     }
   }
 
-  initTrace(initialTrace: Trace) {
+  /**
+   * Initializes the trace with given initialTrace.
+   *
+   * @param {Trace} initialTrace - The initial trace to initialize: the nodes and edges.
+   * @return {TraceableExecution} - The traceable execution object after initialization.
+   */
+  initTrace(initialTrace: Trace): TraceableExecution {
     this.nodes = (initialTrace?.filter((b) => b.group === 'nodes') as Array<Node>) ?? [];
     this.edges = (initialTrace?.filter((b) => b.group === 'edges') as Array<Edge>) ?? [];
     this.narrativesForNonFoundNodes = {};
+    return this;
   }
 
   /**