@outputai/llm 0.13.1-next.83d37f1.0 → 0.13.1-next.a4f6bd4.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@outputai/llm",
-  "version": "0.13.1-next.83d37f1.0",
+  "version": "0.13.1-next.a4f6bd4.0",
   "description": "Framework abstraction to interact with LLM models",
   "type": "module",
   "main": "src/index.js",
@@ -22,7 +22,7 @@
     "gray-matter": "4.0.3",
     "liquidjs": "10.27.2",
     "undici": "8.9.0",
-    "@outputai/core": "0.13.1-next.83d37f1.0"
+    "@outputai/core": "0.13.1-next.a4f6bd4.0"
   },
   "devDependencies": {
     "@ai-sdk/amazon-bedrock": "5.0.57",

package/src/agent.js

@@ -1,6 +1,6 @@
 import { ToolLoopAgent as AIToolLoopAgent } from 'ai';
 import { loadAiSdkTextOptions } from './ai_sdk_options.js';
-import { wrapGeneration, wrapStream } from './utils/wrap.js';
+import { wrapTextGeneration, wrapStream } from './utils/wrap.js';
 import { Role } from './consts.js';
 import { drainStream } from './utils/stream.js';
 import { loadPrompt } from './prompt/loader.js';
@@ -48,15 +48,16 @@ export class Agent {
     const { messages, abortSignal, toolChoice } = Validator.parseAgentGenerateArgs( args );
     const combinedMessages = await this.#combineWithPreviousMessages( messages );
 
-    return wrapGeneration( {
+    return wrapTextGeneration( {
       name: 'Agent.generate',
       prompt: this.#prompt,
-      fn: async () => {
+      fn: async ( { onStepEndHook } ) => {
         const response = await this.#agent.generate( {
           messages: combinedMessages,
           allowSystemInMessages: true,
           ...( abortSignal && { abortSignal } ),
-          ...( toolChoice && { toolChoice } )
+          ...( toolChoice && { toolChoice } ),
+          onStepEnd: onStepEndHook
         } );
         if ( response.finishReason !== 'error' ) {
           await this.#storeMessages( messages.concat( response.responseMessages ?? [] ) );
@@ -73,10 +74,10 @@ export class Agent {
     const { messages, abortSignal, toolChoice, onChunk } = Validator.parseAgentGenerateWithStreamingArgs( args );
     const combinedMessages = await this.#combineWithPreviousMessages( messages );
 
-    return wrapGeneration( {
+    return wrapTextGeneration( {
       name: 'Agent.generateWithStreaming',
       prompt: this.#prompt,
-      fn: async () => {
+      fn: async ( { onStepEndHook } ) => {
         const state = { response: null };
         const stream = await this.#agent.stream( {
           messages: combinedMessages,
@@ -84,6 +85,7 @@ export class Agent {
           ...( onChunk && { onChunk } ),
           ...( abortSignal && { abortSignal } ),
           ...( toolChoice && { toolChoice } ),
+          onStepEnd: onStepEndHook,
           onEnd: res => {
             state.response = res;
           },
@@ -93,7 +95,7 @@ export class Agent {
         await drainStream( stream, abortSignal );
 
         if ( !state.response ) {
-          throw new Error( 'Agent streaming generation completed without a response.' );
+          throw new Error( 'Streaming completed without a response.' );
         }
 
         state.response.output = await stream.output;
@@ -114,20 +116,21 @@ export class Agent {
       name: 'Agent.stream',
       prompt: this.#prompt,
       abortSignal,
-      fn: ( { onEndHook, onErrorHook } ) => this.#agent.stream( {
+      fn: ( { onEndHook, onErrorHook, onStepEndHook, onAbortHook } ) => this.#agent.stream( {
         messages: combinedMessages,
         allowSystemInMessages: true,
         ...( onChunk && { onChunk } ),
         ...( abortSignal && { abortSignal } ),
         ...( toolChoice && { toolChoice } ),
+        onStepEnd: onStepEndHook,
+        onAbort: onAbortHook,
         onEnd: response =>
           onEndHook( response, async parsedResponse => {
             if ( response.finishReason !== 'error' ) {
               await this.#storeMessages( messages.concat( response.responseMessages ?? [] ) )
-                .catch( error => Logger.error( 'Agent.stream message store persistence failed', {
-                  namespace: 'LLM',
-                  error: error instanceof Error ? error.message : String( error )
-                } ) );
+                .catch( error =>
+                  Logger.error( 'Message store persistence failed', { namespace: 'LLM', error: error?.message || String( error ) } )
+                );
             }
             await onEnd?.( parsedResponse );
           } ),

package/src/generate.js

@@ -1,6 +1,6 @@
 import * as AI from 'ai';
 import { loadPrompt } from './prompt/loader.js';
-import { wrapGeneration, wrapStream } from './utils/wrap.js';
+import { wrapImageGeneration, wrapStream, wrapTextGeneration } from './utils/wrap.js';
 import { loadAiSdkTextOptions, loadAiSdkImageOptions } from './ai_sdk_options.js';
 import { drainStream } from './utils/stream.js';
 import { loadSkills } from './utils/skills.js';
@@ -11,10 +11,13 @@ export const generateText = async args => {
   const prompt = promptObject ?? loadPrompt( promptFile, variables, promptDir );
   const skills = loadSkills( prompt );
 
-  return wrapGeneration( {
+  return wrapTextGeneration( {
     name: 'generateText',
     prompt,
-    fn: () => AI.generateText( loadAiSdkTextOptions( { prompt, skills, ...aiOptions } ) )
+    fn: ( { onStepEndHook } ) => AI.generateText( {
+      ...loadAiSdkTextOptions( { prompt, skills, ...aiOptions } ),
+      onStepEnd: onStepEndHook
+    } )
   } );
 };
 
@@ -27,9 +30,11 @@ export const streamText = args => {
     name: 'streamText',
     prompt,
     abortSignal: aiOptions.abortSignal,
-    fn: ( { onEndHook, onErrorHook } ) => AI.streamText( {
+    fn: ( { onEndHook, onErrorHook, onStepEndHook, onAbortHook } ) => AI.streamText( {
       ...loadAiSdkTextOptions( { prompt, skills, ...aiOptions } ),
       ...( onChunk && { onChunk } ),
+      onStepEnd: onStepEndHook,
+      onAbort: onAbortHook,
       onEnd: response => onEndHook( response, onEnd ),
       onError: event => onErrorHook( event, error => onError?.( { ...event, error } ) )
     } )
@@ -44,14 +49,15 @@ export const generateTextWithStreaming = async args => {
   const prompt = promptObject ?? loadPrompt( promptFile, variables, promptDir );
   const skills = loadSkills( prompt );
 
-  return wrapGeneration( {
+  return wrapTextGeneration( {
     name: 'generateTextWithStreaming',
     prompt,
-    fn: async () => {
+    fn: async ( { onStepEndHook } ) => {
       const state = { response: null };
       const stream = AI.streamText( {
         ...loadAiSdkTextOptions( { prompt, skills, ...aiOptions } ),
         ...( onChunk && { onChunk } ),
+        onStepEnd: onStepEndHook,
         onEnd: res => {
           state.response = res;
         },
@@ -61,7 +67,7 @@ export const generateTextWithStreaming = async args => {
       await drainStream( stream, aiOptions.abortSignal );
 
       if ( !state.response ) {
-        throw new Error( 'Streaming generation completed without a response.' );
+        throw new Error( 'Streaming completed without a response.' );
       }
 
       state.response.output = await stream.output;
@@ -74,7 +80,7 @@ export const generateImage = async args => {
   const { promptFile, promptObject, promptDir, variables, ...aiOptions } = Validator.parseGenerateImageArgs( args );
   const prompt = promptObject ?? loadPrompt( promptFile, variables, promptDir );
 
-  return wrapGeneration( {
+  return wrapImageGeneration( {
     name: 'generateImage',
     prompt,
     fn: () => AI.generateImage( loadAiSdkImageOptions( { prompt, ...aiOptions } ) )

package/src/utils/error_handler.js

@@ -25,6 +25,10 @@ const nonRetryableAiSdkErrorTypes = [
  * @returns {object} A new Error
  */
 export const mapAiError = error => {
+  if ( !( error instanceof Error ) ) {
+    return error;
+  }
+
   if ( error instanceof FatalError ) {
     return error;
   }

package/src/utils/image.js

@@ -1,10 +1,37 @@
+import { Logger } from '@outputai/core';
+
 /**
  * Get the approximate file size from a base64 string.
  * @param {string} b64data
  * @returns {number} Size in bytes
  */
 export const calculateBase64FileSize = b64data => {
+  if ( typeof b64data !== 'string' ) {
+    return null;
+  }
   const baseSize = b64data.length * ( 3 / 4 );
   const paddingSize = [ b64data.at( -2 ), b64data.at( -1 ) ].filter( v => v === '=' ).length;
   return baseSize - paddingSize;
 };
+
+/**
+ * Return a serialized version of the images from the AI SDK response.
+ * It contains only the file size and the media-type
+ *
+ * @param {object} response
+ * @returns {Array<{size: number, mediaType: string}>}
+ */
+export const serializeImagesFromResponse = response => {
+  try {
+    if ( !Array.isArray( response?.images ) ) {
+      return [];
+    }
+    return response.images
+      .filter( image => image !== null && typeof image === 'object' )
+      .map( ( { base64, mediaType } ) => ( { size: calculateBase64FileSize( base64 ), mediaType } ) );
+
+  } catch ( error ) {
+    Logger.error( 'Image serialization failed', { namespace: 'LLM', error: error?.message || String( error ) } );
+    return [];
+  }
+};

package/src/utils/metering.js

@@ -0,0 +1,125 @@
+import { parseLLMUsage } from './usage.js';
+import { calculateCosts } from './cost.js';
+import { Tracing, Event } from '@outputai/core/sdk/runtime';
+import { Logger } from '@outputai/core';
+import { convertCostToLegacy } from './legacy_cost_attribute.js';
+
+/**
+ * Collects the usage of a single LLM call and bills it once.
+ *
+ * Steps are recorded as the SDK lifecycle reports them, and the response is handed to `bill()` by
+ * the paths that have one, so a call that throws after the model ran is still billed for what it
+ * spent. `bill()` parses the usage, costs it, attaches both as trace attributes and emits the
+ * metering events; it never throws.
+ */
+export class Metering {
+  #recordedSteps = [];
+  #billed = false;
+  #prompt;
+  #traceId;
+  #attributes = {
+    usage: null,
+    cost: null,
+    legacy: null
+  };
+
+  async #createAttributes( { usage, steps } ) {
+    this.#attributes.usage = parseLLMUsage( { prompt: this.#prompt, usage, steps } );
+    if ( this.#attributes.usage ) {
+      this.#attributes.cost = await calculateCosts( this.#attributes.usage );
+      if ( this.#attributes.cost ) {
+        // @TEMP Preserve the deprecated event and trace attribute for legacy consumers.
+        this.#attributes.legacy = convertCostToLegacy( this.#attributes.cost );
+      }
+    }
+  }
+
+  #attachAttributes() {
+    if ( this.#attributes.usage ) {
+      Tracing.addEventAttribute( { eventId: this.#traceId, attribute: this.#attributes.usage } );
+    }
+    if ( this.#attributes.cost ) {
+      Tracing.addEventAttribute( { eventId: this.#traceId, attribute: this.#attributes.cost } );
+    }
+    if ( this.#attributes.legacy ) {
+      Tracing.addEventAttribute( { eventId: this.#traceId, attribute: this.#attributes.legacy } );
+    }
+  }
+
+  #emitEvents() {
+    if ( this.#attributes.usage ) {
+      Event.emit( 'llm:generation:metering', structuredClone( { cost: this.#attributes.cost, usage: this.#attributes.usage } ) );
+    }
+    if ( this.#attributes.legacy ) {
+      Event.emit( 'cost:llm:request', structuredClone( this.#attributes.legacy ) );
+    }
+  }
+
+  /**
+   * @param {object} args
+   * @param {string} args.traceId - Trace event the usage and cost attributes are attached to
+   * @param {object} args.prompt - Loaded prompt (`config.provider` / `config.model` used for cost)
+   */
+  constructor( { traceId, prompt } ) {
+    this.#traceId = traceId;
+    this.#prompt = prompt;
+  }
+
+  /**
+   * Records a completed step, used as the usage source when the call ends without a response.
+   *
+   * @param {object} step - AI SDK step
+   */
+  recordStep = step => {
+    if ( !this.#billed ) {
+      this.#recordedSteps.push( step );
+    }
+  };
+
+  /**
+   * Bills the recorded usage: parses it, costs it, attaches the trace attributes and emits the
+   * metering events. Later calls and later records are ignored, so it is safe to call from every
+   * path that can end the call. Read the result from `attributes`.
+   *
+   * An empty collector is not a final answer: the SDK reports a stream error before the steps that
+   * preceded it, so billing stays open until there is usage or a step to report.
+   *
+   * @param {object} [response] - AI SDK response, whose aggregate usage and steps take precedence
+   *   over the recorded steps; omitted by the paths that end without one
+   * @returns {Promise<void>}
+   */
+  bill = async ( response = null ) => {
+    if ( this.#billed ) {
+      return;
+    }
+
+    try {
+      const responseSteps = response?.steps;
+      const steps = Array.isArray( responseSteps ) && responseSteps.length > 0 ? responseSteps : this.#recordedSteps;
+      const usage = response?.usage;
+
+      if ( !usage && steps.length === 0 ) {
+        return;
+      }
+
+      this.#billed = true;
+
+      await this.#createAttributes( { usage, steps } );
+      this.#attachAttributes();
+      this.#emitEvents();
+
+    } catch ( error ) {
+      Logger.error( 'Metering failed', { namespace: 'LLM', error: error?.message || String( error ) } );
+    }
+  };
+
+  /**
+   * Attributes produced by `bill()`; each one is null until billed, and stays null when there is
+   * nothing to report.
+   *
+   * @returns {{ usage: object|null, cost: object|null, legacy: object|null }}
+   */
+  get attributes() {
+    return this.#attributes;
+  }
+}