@outputai/llm 0.4.0 → 0.4.1-dev.06c2b50.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@outputai/llm",
-  "version": "0.4.0",
+  "version": "0.4.1-dev.06c2b50.0",
   "description": "Framework abstraction to interact with LLM models",
   "type": "module",
   "main": "src/index.js",
@@ -23,7 +23,7 @@
     "entities": "8.0.0",
     "gray-matter": "4.0.3",
     "liquidjs": "10.25.7",
-    "@outputai/core": "0.4.0"
+    "@outputai/core": "0.4.1-dev.06c2b50.0"
   },
   "license": "Apache-2.0",
   "publishConfig": {

package/src/utils/trace.js

@@ -13,6 +13,8 @@ export const endTraceWithError = ( { traceId, error } ) => {
 export const endTraceWithSuccess = ( { traceId, modelId, response, cost, ...extra } ) => {
   const { totalUsage: usage, text: result, providerMetadata } = response;
   Tracing.addEventAttribute( { eventId: traceId, name: Tracing.Attribute.COST, value: cost } );
-  Tracing.addEventEnd( { id: traceId, details: { result, usage, providerMetadata, ...extra } } );
+  Tracing.addEventAttribute( { eventId: traceId, name: Tracing.Attribute.TOKEN_USAGE, value: usage } );
+  Tracing.addEventEnd( { id: traceId, details: { result, providerMetadata, ...extra } } );
   emitEvent( 'cost:llm:request', { modelId, cost, usage } );
+  emitEvent( 'token_usage:llm:request', { modelId, usage } );
 };

package/src/utils/trace.spec.js

@@ -1,4 +1,6 @@
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import buildTraceTree from '../../../core/src/tracing/tools/build_trace_tree.js';
+import { EventAction } from '../../../core/src/tracing/trace_consts.js';
 
 vi.mock( '@outputai/core/sdk_activity_integration', () => ( {
   Tracing: {
@@ -7,7 +9,8 @@ vi.mock( '@outputai/core/sdk_activity_integration', () => ( {
     addEventAttribute: vi.fn(),
     addEventEnd: vi.fn(),
     Attribute: {
-      COST: 'cost'
+      COST: 'cost',
+      TOKEN_USAGE: 'token_usage'
     }
   },
   emitEvent: vi.fn()
@@ -54,7 +57,7 @@ describe( 'trace utils', () => {
   } );
 
   describe( 'endTraceWithSuccess', () => {
-    it( 'adds cost attribute, ends the trace with response fields and extra details, and emits cost:llm:request', () => {
+    it( 'adds cost + token_usage attributes, ends the trace without usage in output, and emits both cost and token_usage events', () => {
       const cost = { total: 0.01, components: [] };
       const usage = { inputTokens: 2, outputTokens: 3 };
       const response = {
@@ -76,20 +79,63 @@ describe( 'trace utils', () => {
         name: 'cost',
         value: cost
       } );
+      expect( tracing.addEventAttribute ).toHaveBeenCalledWith( {
+        eventId: 'trace-a',
+        name: 'token_usage',
+        value: usage
+      } );
       expect( tracing.addEventEnd ).toHaveBeenCalledWith( {
         id: 'trace-a',
         details: {
           result: 'hello',
-          usage,
           providerMetadata: { provider: 'x' },
           sourcesFromTools: [ { url: 'https://u.test', title: '' } ]
         }
       } );
+      expect( tracing.addEventEnd.mock.calls[0][0].details ).not.toHaveProperty( 'usage' );
       expect( emitEvent ).toHaveBeenCalledWith( 'cost:llm:request', {
         modelId: 'my-model',
         cost,
         usage
       } );
+      expect( emitEvent ).toHaveBeenCalledWith( 'token_usage:llm:request', {
+        modelId: 'my-model',
+        usage
+      } );
+    } );
+
+    it( 'produces an llm trace node with attributes.token_usage and no output.usage when fed through buildTraceTree', () => {
+      const usage = { inputTokens: 12, outputTokens: 7, cachedInputTokens: 3, totalTokens: 22 };
+      const cost = { total: 0.0042, components: [ { name: 'input_tokens', value: 0.002 } ] };
+      const response = {
+        text: 'tree result',
+        totalUsage: usage,
+        providerMetadata: { provider: 'p' }
+      };
+
+      // Capture the calls the wrapper makes against Tracing/emit, and translate them into trace
+      // entries — what buildTraceTree consumes server-side to materialize the persisted trace JSON.
+      endTraceWithSuccess( { traceId: 'llm-1', modelId: 'm', response, cost } );
+
+      const entries = [
+        { kind: 'workflow', action: EventAction.START, name: 'wf', id: 'wf', parentId: undefined, details: {}, timestamp: 1 },
+        { kind: 'llm', action: EventAction.START, name: 'generateText', id: 'llm-1', parentId: 'wf', details: { prompt: 'p' }, timestamp: 10 }
+      ];
+      for ( const call of Tracing.addEventAttribute.mock.calls ) {
+        entries.push( { id: call[0].eventId, action: EventAction.ADD_ATTR, details: { name: call[0].name, value: call[0].value }, timestamp: 20 } );
+      }
+      const endDetails = Tracing.addEventEnd.mock.calls[0][0].details;
+      entries.push( { id: 'llm-1', action: EventAction.END, details: endDetails, timestamp: 30 } );
+      entries.push( { id: 'wf', action: EventAction.END, details: {}, timestamp: 40 } );
+
+      const tree = buildTraceTree( entries );
+      const llmNode = tree.children[0];
+
+      expect( llmNode.kind ).toBe( 'llm' );
+      expect( llmNode.attributes.token_usage ).toEqual( usage );
+      expect( llmNode.attributes.cost ).toEqual( cost );
+      expect( llmNode.output ).not.toHaveProperty( 'usage' );
+      expect( llmNode.output ).toMatchObject( { result: 'tree result', providerMetadata: { provider: 'p' } } );
     } );
   } );
 } );