@outputai/llm 0.4.1-dev.92bc2fb.0 → 0.4.1-next.6bc541c.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@outputai/llm",
-  "version": "0.4.1-dev.92bc2fb.0",
+  "version": "0.4.1-next.6bc541c.0",
   "description": "Framework abstraction to interact with LLM models",
   "type": "module",
   "main": "src/index.js",
@@ -23,7 +23,8 @@
     "entities": "8.0.0",
     "gray-matter": "4.0.3",
     "liquidjs": "10.25.7",
-    "@outputai/core": "0.4.1-dev.92bc2fb.0"
+    "undici": "8.1.0",
+    "@outputai/core": "0.4.1-next.6bc541c.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/utils/trace.js

@@ -13,8 +13,6 @@ 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.addEventAttribute( { eventId: traceId, name: Tracing.Attribute.TOKEN_USAGE, value: usage } );
-  Tracing.addEventEnd( { id: traceId, details: { result, providerMetadata, ...extra } } );
+  Tracing.addEventEnd( { id: traceId, details: { result, usage, providerMetadata, ...extra } } );
   emitEvent( 'cost:llm:request', { modelId, cost, usage } );
-  emitEvent( 'token_usage:llm:request', { modelId, usage } );
 };

package/src/utils/trace.spec.js

@@ -1,6 +1,4 @@
 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: {
@@ -9,8 +7,7 @@ vi.mock( '@outputai/core/sdk_activity_integration', () => ( {
     addEventAttribute: vi.fn(),
     addEventEnd: vi.fn(),
     Attribute: {
-      COST: 'cost',
-      TOKEN_USAGE: 'token_usage'
+      COST: 'cost'
     }
   },
   emitEvent: vi.fn()
@@ -57,7 +54,7 @@ describe( 'trace utils', () => {
   } );
 
   describe( 'endTraceWithSuccess', () => {
-    it( 'adds cost + token_usage attributes, ends the trace without usage in output, and emits both cost and token_usage events', () => {
+    it( 'adds cost attribute, ends the trace with response fields and extra details, and emits cost:llm:request', () => {
       const cost = { total: 0.01, components: [] };
       const usage = { inputTokens: 2, outputTokens: 3 };
       const response = {
@@ -79,63 +76,20 @@ 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' } } );
     } );
   } );
 } );