fix: trace narratives extraction and configuration

tabkram · Nov 26, 2023 · 4806da8 · 4806da8
1 parent 79c5b36
commit 4806da8
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 24 deletions.
diff --git a/src/trace/trace.model.ts b/src/trace/trace.model.ts
@@ -17,7 +17,7 @@ export interface NodeExecutionTraceExtractor<I, O> {
   inputs?: boolean | Array<string> | ((i: I) => unknown);
   outputs?: boolean | Array<string> | ((o: O) => unknown);
   errors?: boolean | Array<string> | ((e: Array<unknown>) => unknown);
-  narratives?: string | ((execTrace: NodeExecutionTrace<I, O>) => Array<string>);
+  narratives?: boolean | Array<string> | ((execTrace: NodeExecutionTrace<I, O>) => Array<string>);
   startTime?: boolean;
   endTime?: boolean;
 }

diff --git a/src/trace/traceableExecution.spec.ts b/src/trace/traceableExecution.spec.ts
@@ -144,30 +144,31 @@ describe('TraceableExecution', () => {
     });
 
     it('should add narratives to a trace node and verify the updated trace and ordered narratives', () => {
-      // Define a sample function
       const sampleFunction = (param: string) => `Result: ${param}`;
 
-      // Run the sample function using the run method and create nodes in the trace
       const nodeId = 'sampleFunction_custom_id_1';
       traceableExecution.run(sampleFunction, ['InputParam'], {
         trace: { id: nodeId, narratives: ['Narrative 0'] },
-        config: {
-          traceExecution: {
-            narratives: (res) => {
-              return [`Narrative 0 with ${res.outputs}`];
-            }
-          }
-        }
+        config: { traceExecution: { narratives: (res) => [`Narrative 0 with ${res.outputs}`] } }
       });
-      traceableExecution.run(sampleFunction, ['InputParam2'], {
-        trace: {
-          id: 'sampleFunction_custom_id_2',
-          narratives: ['Narrative 0 for function 2']
-        }
+
+      traceableExecution.run(sampleFunction, ['InputParam'], {
+        trace: { id: 'sampleFunction_custom_id_2', narratives: ['Narrative 0 for function 2'] },
+        config: { traceExecution: { narratives: ['Narrative 1 for function 2', 'Narrative 2 for function 2'] } }
+      });
+
+      traceableExecution.run(sampleFunction, ['InputParam'], {
+        trace: { id: 'sampleFunction_custom_id_3', narratives: ['Narrative 0 for function 3'] },
+        config: { traceExecution: { narratives: true } }
+      });
+
+      traceableExecution.run(sampleFunction, ['InputParam'], {
+        trace: { id: 'sampleFunction_custom_id_4', narratives: ['Narrative 0 for function 4, non traced!!'] },
+        config: { traceExecution: { narratives: false } }
       });
 
       const trace = traceableExecution.getTrace();
-      expect(trace?.length).toEqual(3);
+      expect(trace?.length).toEqual(7);
 
       // Use pushNarrative to add a single narrative to the specified node
       traceableExecution.pushNarrative(nodeId, 'Narrative 1');
@@ -197,7 +198,10 @@ describe('TraceableExecution', () => {
         'Narrative 1',
         'Narrative 2',
         'Narrative 3',
-        'Narrative 0 for function 2'
+        'Narrative 0 for function 2',
+        'Narrative 1 for function 2',
+        'Narrative 2 for function 2',
+        'Narrative 0 for function 3'
       ]);
     });
 

diff --git a/src/trace/traceableExecution.ts b/src/trace/traceableExecution.ts
@@ -66,10 +66,12 @@ export class TraceableExecution {
     narrativeConfig: NodeExecutionTraceExtractor<I, O>['narratives']
   ): Array<string> {
     try {
-      if (typeof narrativeConfig === 'string') {
-        return [narrativeConfig];
-      } else if (typeof narrativeConfig === 'function') {
-        return narrativeConfig(executionTrace);
+      if (typeof narrativeConfig === 'function') {
+        return (executionTrace.narratives ?? []).concat(narrativeConfig(executionTrace));
+      } else if (Array.isArray(narrativeConfig)) {
+        return (executionTrace.narratives ?? []).concat(narrativeConfig);
+      } else if (narrativeConfig === true) {
+        return executionTrace.narratives;
       }
     } catch (e) {
       throw new Error(`error when mapping/extracting Narrative with config: "${narrativeConfig}", ${e?.message}`);
@@ -452,9 +454,7 @@ export class TraceableExecution {
       execTrace.outputs = TraceableExecution.extractIOExecutionTraceWithConfig<I, O>(executionTrace.outputs, doTraceExecution.outputs);
       execTrace.errors = TraceableExecution.extractIOExecutionTraceWithConfig<I, O>(executionTrace.errors, doTraceExecution.errors);
 
-      execTrace.narratives = (executionTrace.narratives ?? []).concat(
-        TraceableExecution.extractNarrativeWithConfig<I, O>(executionTrace, doTraceExecution.narratives) ?? []
-      );
+      execTrace.narratives = TraceableExecution.extractNarrativeWithConfig<I, O>(executionTrace, doTraceExecution.narratives);
 
       if (doTraceExecution.startTime === true) {
         execTrace.startTime = executionTrace.startTime;