@outputai/core 0.4.1-dev.622e67b.0 → 0.4.1-dev.7aa9a5f.0

package/src/tracing/trace_attribute.js

@@ -0,0 +1,80 @@
+import Decimal from 'decimal.js';
+
+/**
+ * All attributes inherit from this
+ */
+export class BaseAttribute {
+  activityId;
+  activityName;
+  date = Date.now();
+  type;
+
+  constructor( type ) {
+    this.type = type;
+  }
+
+  setActivity( id, name ) {
+    this.activityId = id;
+    this.activityName = name;
+  }
+}
+
+class HTTPRequestCount extends BaseAttribute {
+  static TYPE = 'http:request:count';
+  url;
+  requestId;
+
+  constructor( url, requestId ) {
+    super( HTTPRequestCount.TYPE );
+    this.url = url;
+    this.requestId = requestId;
+  }
+}
+
+class HTTPRequestCost extends BaseAttribute {
+  static TYPE = 'http:request:cost';
+  url;
+  requestId;
+  total = 0;
+
+  constructor( url, requestId, total ) {
+    super( HTTPRequestCost.TYPE );
+    this.url = url;
+    this.requestId = requestId;
+    this.total = total;
+  }
+}
+
+class LLMUsage extends BaseAttribute {
+  static TYPE = 'llm:usage';
+  modelId;
+  usage = [];
+  total = 0;
+  tokensUsed = 0;
+
+  constructor( modelId ) {
+    super( LLMUsage.TYPE );
+    this.modelId = modelId;
+  }
+
+  addUsage( { type, ppm, amount } ) {
+    const total = Decimal( amount ).div( 1_000_000 ).mul( ppm ).toNumber();
+    this.usage.push( {
+      type,
+      ppm,
+      amount,
+      total
+    } );
+    this.total = Decimal( this.total ).add( total ).toNumber();
+    this.tokensUsed = Decimal( this.tokensUsed ).add( amount ).toNumber();
+  }
+}
+
+/**
+ * Types of ADD_ATTR attributes
+ */
+export const Attribute = {
+  LLMUsage,
+  HTTPRequestCost,
+  HTTPRequestCount
+};

package/src/tracing/trace_engine.js

@@ -4,8 +4,10 @@ import { serializeError } from './tools/utils.js';
 import { isStringboolTrue } from '#utils';
 import * as localProcessor from './processors/local/index.js';
 import * as s3Processor from './processors/s3/index.js';
-import { ComponentType } from '#consts';
+import { ComponentType, Signal } from '#consts';
 import { createChildLogger } from '#logger';
+import { EventAction } from './trace_consts.js';
+import { BaseAttribute } from './trace_attribute.js';
 
 const log = createChildLogger( 'Tracing' );
 
@@ -91,7 +93,15 @@ export const addEventAction = ( action, { kind, name, id, parentId, details, exe
 export function addEventActionWithContext( action, options ) {
   const storeContent = Storage.load();
   if ( storeContent ) { // If there is no storageContext this was not called from a Temporal environment
-    const { parentId, executionContext } = storeContent;
+    const { parentId, parentName, executionContext, workflowHandle } = storeContent;
+    if ( action === EventAction.ADD_ATTR ) {
+      const attribute = options.details;
+      if ( !( attribute instanceof BaseAttribute ) ) {
+        throw new Error( `${EventAction.ADD_ATTR} called argument that is not a BaseAttribute instance` );
+      }
+      attribute.setActivity( parentId, parentName );
+      workflowHandle.signal( Signal.ADD_ATTRIBUTE, attribute );
+    }
     addEventAction( action, { ...options, parentId, executionContext } );
   }
 };

package/src/worker/index.js

@@ -69,7 +69,7 @@ const callerDir = process.argv[2];
     workflowsPath,
     activities,
     sinks,
-    interceptors: initInterceptors( { activities, workflows } ),
+    interceptors: initInterceptors( { activities, workflows, connection } ),
     maxConcurrentWorkflowTaskExecutions,
     maxConcurrentActivityTaskExecutions,
     maxCachedWorkflows,

package/src/worker/index.spec.js

@@ -123,7 +123,7 @@ describe( 'worker/index', () => {
       maxConcurrentActivityTaskPolls: configValues.maxConcurrentActivityTaskPolls,
       maxConcurrentWorkflowTaskPolls: configValues.maxConcurrentWorkflowTaskPolls
     } ) );
-    expect( initInterceptorsMock ).toHaveBeenCalledWith( { activities: {}, workflows: [] } );
+    expect( initInterceptorsMock ).toHaveBeenCalledWith( { activities: {}, workflows: [], connection: mockConnection } );
     expect( registerShutdownMock ).toHaveBeenCalledWith( { worker: mockWorker, log: mockLog } );
     expect( startCatalogMock ).toHaveBeenCalledWith( {
       connection: mockConnection,

package/src/worker/interceptors/activity.js

@@ -5,6 +5,7 @@ import { headersToObject } from '../sandboxed_utils.js';
 import { BusEventType, METADATA_ACCESS_SYMBOL } from '#consts';
 import { activityHeartbeatEnabled, activityHeartbeatIntervalMs } from '../configs.js';
 import { messageBus } from '#bus';
+import { Client } from '@temporalio/client';
 
 /*
   This interceptor wraps every activity execution with cross-cutting concerns:
@@ -23,7 +24,7 @@ import { messageBus } from '#bus';
   - Headers injected by the workflow interceptor (executionContext)
 */
 export class ActivityExecutionInterceptor {
-  constructor( { activities, workflows } ) {
+  constructor( { activities, workflows, connection } ) {
     this.activities = activities;
     this.workflowsMap = workflows.reduce( ( map, w ) => {
       map.set( w.name, w );
@@ -32,14 +33,19 @@ export class ActivityExecutionInterceptor {
       }
       return map;
     }, new Map() );
+    this.connection = connection;
   };
 
   async execute( input, next ) {
     const startDate = Date.now();
+    const client = new Client( { connection: this.connection } );
+
     const { workflowExecution: { workflowId }, activityId: id, activityType: name, workflowType: workflowName } = Context.current().info;
     const { executionContext } = headersToObject( input.headers );
     const { type: kind } = this.activities?.[name]?.[METADATA_ACCESS_SYMBOL];
 
+    const workflowHandle = client.workflow.getHandle( workflowId );
+
     messageBus.emit( BusEventType.ACTIVITY_START, { id, name, kind, workflowId, workflowName } );
     Tracing.addEventStart( { id, name, kind, parentId: workflowId, details: input.args[0], executionContext } );
 
@@ -56,7 +62,8 @@ export class ActivityExecutionInterceptor {
       intervals.heartbeat = activityHeartbeatEnabled && setInterval( () => Context.current().heartbeat(), activityHeartbeatIntervalMs );
 
       // Wraps the execution with accessible metadata for the activity
-      const output = await Storage.runWithContext( async _ => next( input ), { parentId: id, executionContext, workflowFilename } );
+      const ctx = { parentId: id, parentName: name, executionContext, workflowFilename, workflowHandle };
+      const output = await Storage.runWithContext( async _ => next( input ), ctx );
 
       messageBus.emit( BusEventType.ACTIVITY_END, { id, name, kind, workflowId, workflowName, duration: Date.now() - startDate } );
       Tracing.addEventEnd( { id, details: output, executionContext } );

package/src/worker/interceptors/activity.spec.js

@@ -2,6 +2,8 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { BusEventType } from '#consts';
 
 const METADATA_ACCESS_SYMBOL = vi.hoisted( () => Symbol( '__metadata' ) );
+const workflowHandleMock = vi.hoisted( () => ( { signal: vi.fn() } ) );
+const getHandleMock = vi.hoisted( () => vi.fn( () => workflowHandleMock ) );
 
 const heartbeatMock = vi.fn();
 const runWithContextMock = vi.hoisted( () => vi.fn().mockImplementation( async fn => fn() ) );
@@ -21,6 +23,14 @@ vi.mock( '@temporalio/activity', () => ( {
   }
 } ) );
 
+vi.mock( '@temporalio/client', () => ( {
+  Client: class Client {
+    workflow = {
+      getHandle: getHandleMock
+    };
+  }
+} ) );
+
 vi.mock( '#async_storage', () => ( {
   Storage: {
     runWithContext: runWithContextMock
@@ -108,12 +118,15 @@ describe( 'ActivityExecutionInterceptor', () => {
     expect( addEventErrorMock ).not.toHaveBeenCalled();
     expect( runWithContextMock ).toHaveBeenCalledWith(
       expect.any( Function ),
-      {
+      expect.objectContaining( {
         parentId: 'act-1',
+        parentName: 'myWorkflow#myStep',
         executionContext: { workflowId: 'wf-1' },
-        workflowFilename: '/workflows/myWorkflow.js'
-      }
+        workflowFilename: '/workflows/myWorkflow.js',
+        workflowHandle: workflowHandleMock
+      } )
     );
+    expect( getHandleMock ).toHaveBeenCalledWith( 'wf-1' );
   } );
 
   it( 'records trace error event on failed execution', async () => {

package/src/worker/interceptors.js

@@ -4,7 +4,7 @@ import { ActivityExecutionInterceptor } from './interceptors/activity.js';
 
 const __dirname = dirname( fileURLToPath( import.meta.url ) );
 
-export const initInterceptors = ( { activities, workflows } ) => ( {
+export const initInterceptors = ( { activities, workflows, connection } ) => ( {
   workflowModules: [ join( __dirname, './interceptors/workflow.js' ) ],
-  activityInbound: [ () => new ActivityExecutionInterceptor( { activities, workflows } ) ]
+  activityInbound: [ () => new ActivityExecutionInterceptor( { activities, workflows, connection } ) ]
 } );