@outputai/llm 0.4.1-next.ae3ab85.0 → 0.4.1-next.d085dde.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@outputai/llm",
-  "version": "0.4.1-next.ae3ab85.0",
+  "version": "0.4.1-next.d085dde.0",
   "description": "Framework abstraction to interact with LLM models",
   "type": "module",
   "main": "src/index.js",
@@ -19,11 +19,11 @@
     "@perplexity-ai/ai-sdk": "0.1.3",
     "@tavily/ai-sdk": "0.4.1",
     "ai": "6.0.168",
-    "decimal.js": "10.6.0",
     "entities": "8.0.0",
     "gray-matter": "4.0.3",
     "liquidjs": "10.25.7",
-    "@outputai/core": "0.4.1-next.ae3ab85.0"
+    "undici": "8.1.0",
+    "@outputai/core": "0.4.1-next.d085dde.0"
   },
   "license": "Apache-2.0",
   "publishConfig": {

package/src/ai_model.js

@@ -1,12 +1,29 @@
-import { bedrock } from '@ai-sdk/amazon-bedrock';
-import { anthropic } from '@ai-sdk/anthropic';
-import { azure } from '@ai-sdk/azure';
-import { vertex } from '@ai-sdk/google-vertex';
-import { openai } from '@ai-sdk/openai';
-import { perplexity } from '@ai-sdk/perplexity';
+import { createAmazonBedrock } from '@ai-sdk/amazon-bedrock';
+import { createAnthropic } from '@ai-sdk/anthropic';
+import { createAzure } from '@ai-sdk/azure';
+import { createVertex } from '@ai-sdk/google-vertex';
+import { createOpenAI } from '@ai-sdk/openai';
+import { createPerplexity } from '@ai-sdk/perplexity';
 import { ValidationError, z } from '@outputai/core';
+import { Agent, fetch } from 'undici';
+
+const dispatcher = new Agent( {
+  headersTimeout: 15 * 60 * 1000, // 15 min
+  bodyTimeout: 15 * 60 * 1000
+} );
+
+const customFetch = ( input, init ) => fetch( input, { dispatcher, ...init } );
+const initProvider = factory => factory( { fetch: customFetch } );
+
+export const builtInProviders = {
+  azure: initProvider( createAzure ),
+  anthropic: initProvider( createAnthropic ),
+  openai: initProvider( createOpenAI ),
+  vertex: initProvider( createVertex ),
+  bedrock: initProvider( createAmazonBedrock ),
+  perplexity: initProvider( createPerplexity )
+};
 
-export const builtInProviders = { azure, anthropic, openai, vertex, bedrock, perplexity };
 export const providers = { ...builtInProviders };
 
 const registerProviderSchema = z.object( {
@@ -14,20 +31,17 @@ const registerProviderSchema = z.object( {
   providerFn: z.function()
 } );
 
+const toolConfigSchema = z.record( z.string(), z.unknown() );
+
 export function registerProvider( name, providerFn ) {
   const result = registerProviderSchema.safeParse( { name, providerFn } );
   if ( !result.success ) {
-    throw new ValidationError(
-      `Invalid provider registration: ${z.prettifyError( result.error )}`,
-      { cause: result.error }
-    );
+    throw new ValidationError( `Invalid provider registration: ${z.prettifyError( result.error )}` );
   }
   providers[name] = providerFn;
 }
 
-export function getRegisteredProviders() {
-  return Object.keys( providers );
-}
+export const getRegisteredProviders = () => Object.keys( providers );
 
 export function loadModel( prompt ) {
   const config = prompt?.config;
@@ -49,10 +63,8 @@ export function loadModel( prompt ) {
   const provider = providers[providerName];
 
   if ( !provider ) {
-    const validProviders = Object.keys( providers ).join( ', ' );
-    throw new Error(
-      `Invalid provider "${providerName}". Valid providers: ${validProviders}`
-    );
+    const availableProviders = Object.keys( providers ).join( ', ' );
+    throw new Error( `Invalid provider "${providerName}". Valid providers: ${availableProviders}` );
   }
 
   return provider( modelName );
@@ -60,7 +72,12 @@ export function loadModel( prompt ) {
 
 export function loadTools( prompt ) {
   const config = prompt?.config;
-  const toolsConfig = config?.tools;
+
+  if ( !config ) {
+    throw new Error( 'Prompt is missing config object' );
+  }
+
+  const { tools: toolsConfig, provider: providerName } = config;
 
   if ( !toolsConfig ) {
     return null;
@@ -84,20 +101,15 @@ export function loadTools( prompt ) {
     return null;
   }
 
-  const providerName = config.provider;
   const provider = providers[providerName];
 
   if ( !provider ) {
-    const validProviders = Object.keys( providers ).join( ', ' );
-    throw new Error(
-      `Invalid provider "${providerName}". Valid providers: ${validProviders}`
-    );
+    const availableProviders = Object.keys( providers ).join( ', ' );
+    throw new Error( `Invalid provider "${providerName}". Valid providers: ${availableProviders}` );
   }
 
   if ( !provider.tools || typeof provider.tools !== 'object' ) {
-    throw new Error(
-      `Provider "${providerName}" does not support provider-specific tools.`
-    );
+    throw new Error( `Provider "${providerName}" does not support provider-specific tools.` );
   }
 
   const tools = {};
@@ -109,18 +121,14 @@ export function loadTools( prompt ) {
       const availableTools = Object.keys( provider.tools )
         .filter( key => typeof provider.tools[key] === 'function' )
         .join( ', ' );
+      const toolsMessage = availableTools ? `Available tools: ${availableTools}` : 'No tools are available';
 
-      throw new Error(
-        `Unknown tool "${toolName}" for provider "${providerName}".` +
-        ( availableTools ? ` Available tools: ${availableTools}` : '' )
-      );
+      throw new Error( `Unknown tool "${toolName}" for provider "${providerName}". ${toolsMessage}` );
     }
 
-    if ( typeof toolConfig !== 'object' || toolConfig === null ) {
-      throw new Error(
-        `Configuration for tool "${toolName}" must be an object. ` +
-        `Use "${toolName}: {}" for tools without configuration.`
-      );
+    const result = toolConfigSchema.safeParse( toolConfig );
+    if ( !result.success ) {
+      throw new ValidationError( `Invalid config for tool "${toolName}": ${z.prettifyError( result.error )}` );
     }
 
     tools[toolName] = toolFactory( toolConfig );

package/src/ai_model.spec.js

@@ -1,10 +1,12 @@
 import { it, expect, vi, afterEach, describe } from 'vitest';
 
-const openaiImpl = vi.fn( model => `openai:${model}` );
-const azureImpl = vi.fn( model => `azure:${model}` );
-const anthropicImpl = vi.fn( model => `anthropic:${model}` );
-const bedrockImpl = vi.fn( model => `bedrock:${model}` );
-const perplexityImpl = vi.fn( model => `perplexity:${model}` );
+const providerFactoryOptions = vi.hoisted( () => ( {} ) );
+const openaiImpl = vi.hoisted( () => vi.fn( model => `openai:${model}` ) );
+const azureImpl = vi.hoisted( () => vi.fn( model => `azure:${model}` ) );
+const anthropicImpl = vi.hoisted( () => vi.fn( model => `anthropic:${model}` ) );
+const bedrockImpl = vi.hoisted( () => vi.fn( model => `bedrock:${model}` ) );
+const perplexityImpl = vi.hoisted( () => vi.fn( model => `perplexity:${model}` ) );
+const vertexImpl = vi.hoisted( () => vi.fn( model => `vertex:${model}` ) );
 
 // OpenAI mock with tools support
 vi.mock( '@ai-sdk/openai', () => {
@@ -12,12 +14,20 @@ vi.mock( '@ai-sdk/openai', () => {
   openaiMock.tools = {
     webSearch: ( config = {} ) => ( { type: 'webSearch', config } )
   };
-  return { openai: openaiMock };
+  return {
+    createOpenAI: options => {
+      providerFactoryOptions.openai = options;
+      return openaiMock;
+    }
+  };
 } );
 
 // Azure mock without tools support
 vi.mock( '@ai-sdk/azure', () => ( {
-  azure: ( ...values ) => azureImpl( ...values )
+  createAzure: options => {
+    providerFactoryOptions.azure = options;
+    return ( ...values ) => azureImpl( ...values );
+  }
 } ) );
 
 // Anthropic mock with tools support
@@ -30,7 +40,12 @@ vi.mock( '@ai-sdk/anthropic', () => {
     codeExecution_20250522: ( config = {} ) => ( { type: 'codeExecution_20250522', config } ),
     codeExecution_20250825: ( config = {} ) => ( { type: 'codeExecution_20250825', config } )
   };
-  return { anthropic: anthropicMock };
+  return {
+    createAnthropic: options => {
+      providerFactoryOptions.anthropic = options;
+      return anthropicMock;
+    }
+  };
 } );
 
 // Bedrock mock with tools support
@@ -42,17 +57,25 @@ vi.mock( '@ai-sdk/amazon-bedrock', () => {
     textEditor_20250429: ( config = {} ) => ( { type: 'textEditor_20250429', config } ),
     computer_20241022: ( config = {} ) => ( { type: 'computer_20241022', config } )
   };
-  return { bedrock: bedrockMock };
+  return {
+    createAmazonBedrock: options => {
+      providerFactoryOptions.bedrock = options;
+      return bedrockMock;
+    }
+  };
 } );
 
 // Perplexity mock
 vi.mock( '@ai-sdk/perplexity', () => ( {
-  perplexity: ( ...values ) => perplexityImpl( ...values )
+  createPerplexity: options => {
+    providerFactoryOptions.perplexity = options;
+    return ( ...values ) => perplexityImpl( ...values );
+  }
 } ) );
 
 // Vertex mock with tools support
 vi.mock( '@ai-sdk/google-vertex', () => {
-  const vertexFn = model => `vertex:${model}`;
+  const vertexFn = ( ...values ) => vertexImpl( ...values );
   vertexFn.tools = {
     googleSearch: ( config = {} ) => ( { type: 'googleSearch', config } ),
     fileSearch: ( config = {} ) => ( { type: 'fileSearch', config } ),
@@ -62,7 +85,12 @@ vi.mock( '@ai-sdk/google-vertex', () => {
     codeExecution: ( config = {} ) => ( { type: 'codeExecution', config } ),
     vertexRagStore: ( config = {} ) => ( { type: 'vertexRagStore', config } )
   };
-  return { vertex: vertexFn };
+  return {
+    createVertex: options => {
+      providerFactoryOptions.vertex = options;
+      return vertexFn;
+    }
+  };
 } );
 
 import { loadModel, loadTools, registerProvider, getRegisteredProviders, providers, builtInProviders } from './ai_model.js';
@@ -73,6 +101,17 @@ afterEach( async () => {
 } );
 
 describe( 'loadModel', () => {
+  it( 'initializes built-in providers with custom fetch', () => {
+    expect( providerFactoryOptions ).toMatchObject( {
+      azure: { fetch: expect.any( Function ) },
+      anthropic: { fetch: expect.any( Function ) },
+      openai: { fetch: expect.any( Function ) },
+      vertex: { fetch: expect.any( Function ) },
+      bedrock: { fetch: expect.any( Function ) },
+      perplexity: { fetch: expect.any( Function ) }
+    } );
+  } );
+
   it( 'loads model using selected provider', () => {
     const result = loadModel( { config: { provider: 'openai', model: 'gpt-4o-mini' } } );
 
@@ -617,7 +656,7 @@ describe( 'loadTools', () => {
           provider: 'vertex',
           tools: { googleSearch: null }
         }
-      } ) ).toThrow( 'Configuration for tool "googleSearch" must be an object' );
+      } ) ).toThrow( /Invalid config for tool "googleSearch".*expected record, received null/s );
     } );
 
     it( 'throws error when tool config is a string', () => {
@@ -626,7 +665,7 @@ describe( 'loadTools', () => {
           provider: 'vertex',
           tools: { googleSearch: 'MODE_DYNAMIC' }
         }
-      } ) ).toThrow( 'Configuration for tool "googleSearch" must be an object' );
+      } ) ).toThrow( /Invalid config for tool "googleSearch".*expected record, received string/s );
     } );
 
     it( 'throws error for unknown tool on Bedrock with dynamic tool listing', () => {

package/src/ai_sdk.js

@@ -64,10 +64,10 @@ export const hydratePromptTemplate = ( prompt, variables, promptDir, callerSkill
   return { loadedPrompt: meta, allVariables: variables, tools };
 };
 
-export async function generateText( { prompt, variables, promptDir, skills = [], maxSteps = 10, ...extraAiSdkOptions } ) {
+export async function generateText( { prompt, variables, promptDir, skills = [], maxSteps = 10, ...aiSdkArgs } ) {
   const callerSkills = typeof skills === 'function' ? await skills( variables ) : skills;
   const { loadedPrompt, allVariables, tools } =
-    hydratePromptTemplate( prompt, variables, promptDir, callerSkills, extraAiSdkOptions.tools );
+    hydratePromptTemplate( prompt, variables, promptDir, callerSkills, aiSdkArgs.tools );
   const hasTools = Object.keys( tools ).length > 0;
 
   validateGenerateTextArgs( { prompt, variables: allVariables } );
@@ -78,9 +78,10 @@ export async function generateText( { prompt, variables, promptDir, skills = [],
   try {
     const response = await AI.generateText( {
       ...loadAiSdkOptionsFromPrompt( loadedPrompt ),
-      ...extraAiSdkOptions,
+      maxRetries: 0,
+      ...aiSdkArgs,
       ...( hasTools ? { tools } : {} ),
-      ...( hasTools && !extraAiSdkOptions.stopWhen ? { stopWhen: stepCountIs( maxSteps ) } : {} )
+      ...( hasTools && !aiSdkArgs.stopWhen ? { stopWhen: stepCountIs( maxSteps ) } : {} )
     } );
     return wrapTextResponse( { traceId, modelId, response } );
   } catch ( error ) {
@@ -89,7 +90,7 @@ export async function generateText( { prompt, variables, promptDir, skills = [],
   }
 }
 
-export function streamText( { prompt, variables, onFinish, onError, ...restOptions } ) {
+export function streamText( { prompt, variables, onFinish, onError, ...aiSdkArgs } ) {
   validateStreamTextArgs( { prompt, variables } );
   const loadedPrompt = loadPrompt( prompt, variables );
   const traceId = startTrace( { name: 'streamText', prompt, variables, loadedPrompt } );
@@ -98,7 +99,8 @@ export function streamText( { prompt, variables, onFinish, onError, ...restOptio
   try {
     return AI.streamText( {
       ...loadAiSdkOptionsFromPrompt( loadedPrompt ),
-      ...restOptions,
+      maxRetries: 0,
+      ...aiSdkArgs,
       ...wrapStreamOnFinishResponse( { traceId, modelId, onFinish } ),
       onError( event ) {
         endTraceWithError( { traceId, error: event.error } );

package/src/ai_sdk.spec.js

@@ -149,6 +149,7 @@ describe( 'ai_sdk', () => {
       model: 'MODEL',
       messages: basePrompt.messages,
       temperature: 0.3,
+      maxRetries: 0,
       providerOptions: basePrompt.config.providerOptions
     } );
     expect( result ).toBe( generateTextAiFixture.response );
@@ -182,6 +183,7 @@ describe( 'ai_sdk', () => {
     expect( aiFns.generateText ).toHaveBeenCalledWith( {
       model: 'MODEL',
       messages: promptWithProviderOptions.messages,
+      maxRetries: 0,
       providerOptions: {
         thinking: {
           type: 'enabled',
@@ -362,6 +364,7 @@ describe( 'ai_sdk', () => {
         model: 'MODEL',
         messages: basePrompt.messages,
         temperature: 0.3,
+        maxRetries: 0,
         providerOptions: basePrompt.config.providerOptions,
         onFinish: expect.any( Function ),
         onError: expect.any( Function )

package/src/cost/index.js

@@ -1,8 +1,5 @@
 import { fetchModelsPricing } from './fetch_models_pricing.js';
-import Decimal from 'decimal.js';
-
-const M = 1_000_000;
-const calcCost = ( tokens, ppm ) => Decimal( tokens ?? 0 ).div( M ).mul( ppm ).toNumber();
+import { Tracing } from '@outputai/core/sdk_activity_integration';
 
 /**
  * Calculates the cost of an llm call based on the model and usage.
@@ -14,29 +11,41 @@ const calcCost = ( tokens, ppm ) => Decimal( tokens ?? 0 ).div( M ).mul( ppm ).t
 export const calculateLLMCallCost = async ( { modelId, usage } ) => {
   try {
     const models = await fetchModelsPricing();
+
     if ( !models ) {
-      return { total: null, message: 'Failed to fetch models pricing' };
+      console.warn( 'Failed to fetch models pricing' );
+      return null;
     }
 
-    const cost = models.get( modelId );
-    if ( !cost ) {
-      return { total: null, message: 'Missing cost reference for model' };
+    const pricing = models.get( modelId );
+    if ( !pricing ) {
+      console.warn( 'Missing cost reference for model' );
+      return null;
     }
 
     const { inputTokens, cachedInputTokens, outputTokens, reasoningTokens } = usage;
 
     const nonCachedTokens = inputTokens - ( cachedInputTokens ?? 0 );
 
-    const components = [
-      Number.isFinite( cost.input ) ? { name: 'input_tokens', value: calcCost( nonCachedTokens, cost.input ) } : false,
-      Number.isFinite( cost.cache_read ) ? { name: 'input_cached_tokens', value: calcCost( cachedInputTokens, cost.cache_read ) } : false,
-      Number.isFinite( cost.output ) ? { name: 'output_tokens', value: calcCost( outputTokens, cost.output ) } : false,
-      /* When there aren't reasoning costs, the providers doesn't differentiate reasoning vs output, so the price is included in the output */
-      Number.isFinite( cost.reasoning ) ? { name: 'reasoning_tokens', value: calcCost( reasoningTokens, cost.reasoning ) } : false
-    ].filter( v => !!v );
-    return { total: components.reduce( ( v, e ) => v.plus( e.value ), Decimal( 0 ) ).toNumber(), components };
+    const llmUsage = new Tracing.Attribute.LLMUsage( modelId );
+
+    if ( Number.isFinite( pricing.input ) && Number.isFinite( nonCachedTokens ) ) {
+      llmUsage.addUsage( { type: 'input', ppm: pricing.input, amount: nonCachedTokens } );
+    }
+    if ( Number.isFinite( pricing.cache_read ) && Number.isFinite( cachedInputTokens ) ) {
+      llmUsage.addUsage( { type: 'input_cached', ppm: pricing.cache_read, amount: cachedInputTokens } );
+    }
+    if ( Number.isFinite( pricing.output ) && Number.isFinite( outputTokens ) ) {
+      llmUsage.addUsage( { type: 'output', ppm: pricing.output, amount: outputTokens } );
+    }
+    // When there aren't reasoning costs, the providers doesn't differentiate reasoning vs output, so the price is included in the output
+    if ( Number.isFinite( pricing.reasoning ) && Number.isFinite( reasoningTokens ) ) {
+      llmUsage.addUsage( { type: 'reasoning', ppm: pricing.reasoning, amount: reasoningTokens } );
+    }
+
+    return llmUsage;
   } catch ( error ) {
     console.error( 'Error calculating LLM call costs', error );
-    return { total: null, message: `Error calculating LLM call costs: ${error.constructor.name} - ${error.message}` };
+    return null;
   }
 };

package/src/cost/index.spec.js

@@ -1,18 +1,75 @@
-import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+
+const mockFetchModelsPricing = vi.hoisted( () => vi.fn() );
 
-const mockFetchModelsPricing = vi.fn();
 vi.mock( './fetch_models_pricing.js', () => ( {
   fetchModelsPricing: ( ...args ) => mockFetchModelsPricing( ...args )
 } ) );
 
+vi.mock( '@outputai/core/sdk_activity_integration', () => {
+  class LLMUsage {
+    static TYPE = 'llm:usage';
+    type = LLMUsage.TYPE;
+    modelId;
+    usage = [];
+
+    constructor( modelId ) {
+      this.modelId = modelId;
+    }
+
+    addUsage( { type, ppm, amount } ) {
+      this.usage.push( {
+        type,
+        ppm,
+        amount,
+        total: ( amount / 1_000_000 ) * ppm
+      } );
+    }
+
+    get total() {
+      return this.usage.reduce( ( total, current ) => total + current.total, 0 );
+    }
+
+    get tokensUsed() {
+      return this.usage.reduce( ( total, current ) => total + current.amount, 0 );
+    }
+  }
+
+  return {
+    Tracing: {
+      Attribute: {
+        LLMUsage
+      }
+    }
+  };
+} );
+
+import { Tracing } from '@outputai/core/sdk_activity_integration';
 import { calculateLLMCallCost } from './index.js';
 
+const expectLLMUsage = ( result, { modelId, usage, total, tokensUsed } ) => {
+  expect( result ).toBeInstanceOf( Tracing.Attribute.LLMUsage );
+  expect( result ).toEqual( expect.objectContaining( {
+    type: Tracing.Attribute.LLMUsage.TYPE,
+    modelId,
+    usage
+  } ) );
+  expect( result.total ).toBeCloseTo( total );
+  expect( result.tokensUsed ).toBe( tokensUsed );
+};
+
 describe( 'calculateLLMCallCost', () => {
   beforeEach( () => {
     vi.clearAllMocks();
+    vi.spyOn( console, 'warn' ).mockImplementation( () => {} );
+    vi.spyOn( console, 'error' ).mockImplementation( () => {} );
   } );
 
-  it( 'returns total null and message when fetchModelsPricing returns null', async () => {
+  afterEach( () => {
+    vi.restoreAllMocks();
+  } );
+
+  it( 'returns null when fetchModelsPricing returns null', async () => {
     mockFetchModelsPricing.mockResolvedValue( null );
 
     const result = await calculateLLMCallCost( {
@@ -20,10 +77,11 @@ describe( 'calculateLLMCallCost', () => {
       usage: { inputTokens: 100, outputTokens: 50 }
     } );
 
-    expect( result ).toEqual( { total: null, message: 'Failed to fetch models pricing' } );
+    expect( result ).toBeNull();
+    expect( console.warn ).toHaveBeenCalledWith( 'Failed to fetch models pricing' );
   } );
 
-  it( 'returns total null and message when model is missing from cost table', async () => {
+  it( 'returns null when model is missing from cost table', async () => {
     mockFetchModelsPricing.mockResolvedValue( new Map() );
 
     const result = await calculateLLMCallCost( {
@@ -31,47 +89,50 @@ describe( 'calculateLLMCallCost', () => {
       usage: { inputTokens: 100, outputTokens: 50 }
     } );
 
-    expect( result ).toEqual( {
-      total: null,
-      message: 'Missing cost reference for model'
-    } );
+    expect( result ).toBeNull();
+    expect( console.warn ).toHaveBeenCalledWith( 'Missing cost reference for model' );
   } );
 
-  it( 'calculates input and output cost from mock model', async () => {
-    const cost = { input: 2, output: 10, cache_read: 1 };
-    mockFetchModelsPricing.mockResolvedValue( new Map( [ [ 'gpt-4o', cost ] ] ) );
+  it( 'calculates input and output usage from model pricing', async () => {
+    mockFetchModelsPricing.mockResolvedValue( new Map( [ [ 'gpt-4o', { input: 2, output: 10, cache_read: 1 } ] ] ) );
 
     const result = await calculateLLMCallCost( {
       modelId: 'gpt-4o',
       usage: { inputTokens: 1_000_000, outputTokens: 500_000 }
     } );
 
-    expect( result.total ).toBe( 7 );
-    expect( result.components ).toEqual( [
-      { name: 'input_tokens', value: 2 },
-      { name: 'input_cached_tokens', value: 0 },
-      { name: 'output_tokens', value: 5 }
-    ] );
+    expectLLMUsage( result, {
+      modelId: 'gpt-4o',
+      usage: [
+        { type: 'input', ppm: 2, amount: 1_000_000, total: 2 },
+        { type: 'output', ppm: 10, amount: 500_000, total: 5 }
+      ],
+      total: 7,
+      tokensUsed: 1_500_000
+    } );
   } );
 
-  it( 'splits input into non-cached and cached at respective rates', async () => {
-    const cost = { input: 4, cache_read: 1, output: 10 };
-    mockFetchModelsPricing.mockResolvedValue( new Map( [ [ 'cached-model', cost ] ] ) );
+  it( 'splits input into non-cached and cached usage at respective rates', async () => {
+    mockFetchModelsPricing.mockResolvedValue( new Map( [ [ 'cached-model', { input: 4, cache_read: 1, output: 10 } ] ] ) );
 
     const result = await calculateLLMCallCost( {
       modelId: 'cached-model',
       usage: { inputTokens: 1_000_000, cachedInputTokens: 500_000, outputTokens: 100_000 }
     } );
 
-    expect( result.components ).toEqual( [
-      { name: 'input_tokens', value: 2 },
-      { name: 'input_cached_tokens', value: 0.5 },
-      { name: 'output_tokens', value: 1 }
-    ] );
-    expect( result.total ).toBeCloseTo( 3.5 );
+    expectLLMUsage( result, {
+      modelId: 'cached-model',
+      usage: [
+        { type: 'input', ppm: 4, amount: 500_000, total: 2 },
+        { type: 'input_cached', ppm: 1, amount: 500_000, total: 0.5 },
+        { type: 'output', ppm: 10, amount: 100_000, total: 1 }
+      ],
+      total: 3.5,
+      tokensUsed: 1_100_000
+    } );
   } );
 
-  it( 'omits cached component when model has no cache_read (non-cached rate applies to full input minus cached)', async () => {
+  it( 'omits cached usage when model has no cache_read rate', async () => {
     mockFetchModelsPricing.mockResolvedValue( new Map( [ [ 'no-cache', { input: 2, output: 10 } ] ] ) );
 
     const result = await calculateLLMCallCost( {
@@ -79,14 +140,18 @@ describe( 'calculateLLMCallCost', () => {
       usage: { inputTokens: 1_000_000, cachedInputTokens: 200_000, outputTokens: 0 }
     } );
 
-    expect( result.components ).toEqual( [
-      { name: 'input_tokens', value: 1.6 },
-      { name: 'output_tokens', value: 0 }
-    ] );
-    expect( result.total ).toBe( 1.6 );
+    expectLLMUsage( result, {
+      modelId: 'no-cache',
+      usage: [
+        { type: 'input', ppm: 2, amount: 800_000, total: 1.6 },
+        { type: 'output', ppm: 10, amount: 0, total: 0 }
+      ],
+      total: 1.6,
+      tokensUsed: 800_000
+    } );
   } );
 
-  it( 'omits input component when pricing has no input rate', async () => {
+  it( 'omits input usage when pricing has no input rate', async () => {
     mockFetchModelsPricing.mockResolvedValue( new Map( [ [ 'out-only', { output: 10 } ] ] ) );
 
     const result = await calculateLLMCallCost( {
@@ -94,13 +159,17 @@ describe( 'calculateLLMCallCost', () => {
       usage: { inputTokens: 100, outputTokens: 50 }
     } );
 
-    expect( result.total ).toBe( 0.0005 );
-    expect( result.components ).toEqual( [
-      { name: 'output_tokens', value: 0.0005 }
-    ] );
+    expectLLMUsage( result, {
+      modelId: 'out-only',
+      usage: [
+        { type: 'output', ppm: 10, amount: 50, total: 0.0005 }
+      ],
+      total: 0.0005,
+      tokensUsed: 50
+    } );
   } );
 
-  it( 'omits output component when pricing has no output rate', async () => {
+  it( 'omits output usage when pricing has no output rate', async () => {
     mockFetchModelsPricing.mockResolvedValue( new Map( [ [ 'in-only', { input: 1 } ] ] ) );
 
     const result = await calculateLLMCallCost( {
@@ -108,13 +177,17 @@ describe( 'calculateLLMCallCost', () => {
       usage: { inputTokens: 100, outputTokens: 50 }
     } );
 
-    expect( result.total ).toBe( 0.0001 );
-    expect( result.components ).toEqual( [
-      { name: 'input_tokens', value: 0.0001 }
-    ] );
+    expectLLMUsage( result, {
+      modelId: 'in-only',
+      usage: [
+        { type: 'input', ppm: 1, amount: 100, total: 0.0001 }
+      ],
+      total: 0.0001,
+      tokensUsed: 100
+    } );
   } );
 
-  it( 'uses reasoning cost when present', async () => {
+  it( 'includes reasoning usage when present', async () => {
     mockFetchModelsPricing.mockResolvedValue( new Map( [ [
       'with-reasoning',
       { input: 1, output: 10, reasoning: 60 }
@@ -125,15 +198,19 @@ describe( 'calculateLLMCallCost', () => {
       usage: { inputTokens: 100, outputTokens: 20, reasoningTokens: 50 }
     } );
 
-    expect( result.total ).toBeCloseTo( 0.0033 );
-    expect( result.components ).toEqual( [
-      { name: 'input_tokens', value: 0.0001 },
-      { name: 'output_tokens', value: 0.0002 },
-      { name: 'reasoning_tokens', value: 0.003 }
-    ] );
+    expectLLMUsage( result, {
+      modelId: 'with-reasoning',
+      usage: [
+        { type: 'input', ppm: 1, amount: 100, total: 0.0001 },
+        { type: 'output', ppm: 10, amount: 20, total: 0.0002 },
+        { type: 'reasoning', ppm: 60, amount: 50, total: 0.003 }
+      ],
+      total: 0.0033,
+      tokensUsed: 170
+    } );
   } );
 
-  it( 'omits reasoning component when reasoning cost missing (included in output)', async () => {
+  it( 'omits reasoning usage when reasoning cost is missing', async () => {
     mockFetchModelsPricing.mockResolvedValue( new Map( [ [ 'no-reasoning', { input: 1, output: 10 } ] ] ) );
 
     const result = await calculateLLMCallCost( {
@@ -141,14 +218,18 @@ describe( 'calculateLLMCallCost', () => {
       usage: { inputTokens: 100, outputTokens: 20, reasoningTokens: 50 }
     } );
 
-    expect( result.total ).toBeCloseTo( 0.0003 );
-    expect( result.components ).toEqual( [
-      { name: 'input_tokens', value: 0.0001 },
-      { name: 'output_tokens', value: 0.0002 }
-    ] );
+    expectLLMUsage( result, {
+      modelId: 'no-reasoning',
+      usage: [
+        { type: 'input', ppm: 1, amount: 100, total: 0.0001 },
+        { type: 'output', ppm: 10, amount: 20, total: 0.0002 }
+      ],
+      total: 0.0003,
+      tokensUsed: 120
+    } );
   } );
 
-  it( 'includes reasoning component with zero when reasoningTokens is zero', async () => {
+  it( 'includes reasoning usage with zero amount when reasoningTokens is zero', async () => {
     mockFetchModelsPricing.mockResolvedValue( new Map( [ [
       'full',
       { input: 2, output: 8, reasoning: 60 }
@@ -159,15 +240,19 @@ describe( 'calculateLLMCallCost', () => {
       usage: { inputTokens: 100, outputTokens: 50, reasoningTokens: 0 }
     } );
 
-    expect( result.components ).toEqual( [
-      { name: 'input_tokens', value: 0.0002 },
-      { name: 'output_tokens', value: 0.0004 },
-      { name: 'reasoning_tokens', value: 0 }
-    ] );
-    expect( result.total ).toBeCloseTo( 0.0006 );
+    expectLLMUsage( result, {
+      modelId: 'full',
+      usage: [
+        { type: 'input', ppm: 2, amount: 100, total: 0.0002 },
+        { type: 'output', ppm: 8, amount: 50, total: 0.0004 },
+        { type: 'reasoning', ppm: 60, amount: 0, total: 0 }
+      ],
+      total: 0.0006,
+      tokensUsed: 150
+    } );
   } );
 
-  it( 'treats null/undefined token counts as 0', async () => {
+  it( 'omits usage entries for non-finite token counts', async () => {
     mockFetchModelsPricing.mockResolvedValue( new Map( [ [ 'm', { input: 1, output: 2 } ] ] ) );
 
     const result = await calculateLLMCallCost( {
@@ -175,10 +260,26 @@ describe( 'calculateLLMCallCost', () => {
       usage: { inputTokens: null, outputTokens: undefined }
     } );
 
-    expect( result.total ).toBe( 0 );
-    expect( result.components ).toEqual( [
-      { name: 'input_tokens', value: 0 },
-      { name: 'output_tokens', value: 0 }
-    ] );
+    expectLLMUsage( result, {
+      modelId: 'm',
+      usage: [
+        { type: 'input', ppm: 1, amount: 0, total: 0 }
+      ],
+      total: 0,
+      tokensUsed: 0
+    } );
+  } );
+
+  it( 'returns null when pricing lookup throws', async () => {
+    const error = new Error( 'boom' );
+    mockFetchModelsPricing.mockRejectedValue( error );
+
+    const result = await calculateLLMCallCost( {
+      modelId: 'gpt-4o',
+      usage: { inputTokens: 100, outputTokens: 50 }
+    } );
+
+    expect( result ).toBeNull();
+    expect( console.error ).toHaveBeenCalledWith( 'Error calculating LLM call costs', error );
   } );
 } );

package/src/utils/trace.js

@@ -12,7 +12,9 @@ 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 } );
+  if ( cost ) {
+    Tracing.addEventAttribute( { eventId: traceId, attribute: cost } );
+    emitEvent( 'cost:llm:request', cost );
+  }
   Tracing.addEventEnd( { id: traceId, details: { result, usage, providerMetadata, ...extra } } );
-  emitEvent( 'cost:llm:request', { modelId, cost, usage } );
 };

package/src/utils/trace.spec.js

@@ -5,10 +5,7 @@ vi.mock( '@outputai/core/sdk_activity_integration', () => ( {
     addEventStart: vi.fn(),
     addEventError: vi.fn(),
     addEventAttribute: vi.fn(),
-    addEventEnd: vi.fn(),
-    Attribute: {
-      COST: 'cost'
-    }
+    addEventEnd: vi.fn()
   },
   emitEvent: vi.fn()
 } ) );
@@ -54,8 +51,8 @@ describe( 'trace utils', () => {
   } );
 
   describe( 'endTraceWithSuccess', () => {
-    it( 'adds cost attribute, ends the trace with response fields and extra details, and emits cost:llm:request', () => {
-      const cost = { total: 0.01, components: [] };
+    it( 'adds cost attribute, emits cost attribute, and ends the trace with response fields and extra details', () => {
+      const cost = { type: 'llm:usage', modelId: 'my-model', total: 0.01, usage: [] };
       const usage = { inputTokens: 2, outputTokens: 3 };
       const response = {
         text: 'hello',
@@ -73,9 +70,9 @@ describe( 'trace utils', () => {
 
       expect( tracing.addEventAttribute ).toHaveBeenCalledWith( {
         eventId: 'trace-a',
-        name: 'cost',
-        value: cost
+        attribute: cost
       } );
+      expect( emitEvent ).toHaveBeenCalledWith( 'cost:llm:request', cost );
       expect( tracing.addEventEnd ).toHaveBeenCalledWith( {
         id: 'trace-a',
         details: {
@@ -85,10 +82,31 @@ describe( 'trace utils', () => {
           sourcesFromTools: [ { url: 'https://u.test', title: '' } ]
         }
       } );
-      expect( emitEvent ).toHaveBeenCalledWith( 'cost:llm:request', {
+    } );
+
+    it( 'does not emit or add an attribute when cost is missing', () => {
+      const usage = { inputTokens: 2, outputTokens: 3 };
+      const response = {
+        text: 'hello',
+        totalUsage: usage,
+        providerMetadata: { provider: 'x' }
+      };
+
+      endTraceWithSuccess( {
+        traceId: 'trace-no-cost',
         modelId: 'my-model',
-        cost,
-        usage
+        response
+      } );
+
+      expect( tracing.addEventAttribute ).not.toHaveBeenCalled();
+      expect( emitEvent ).not.toHaveBeenCalled();
+      expect( tracing.addEventEnd ).toHaveBeenCalledWith( {
+        id: 'trace-no-cost',
+        details: {
+          result: 'hello',
+          usage,
+          providerMetadata: { provider: 'x' }
+        }
       } );
     } );
   } );