@outputai/llm 0.6.1-dev.daae905.0 → 0.6.1-next.2cc4685.0

package/src/ai_model.js

@@ -31,8 +31,45 @@ const registerProviderSchema = z.object( {
   providerFn: z.function()
 } );
 
+const modelPromptSchema = z.object( {
+  config: z.object( {
+    provider: z.string().min( 1 ),
+    model: z.string().min( 1 )
+  } ).loose()
+} ).loose();
+
+const toolsPromptSchema = z.object( {
+  config: z.object( {
+    provider: z.string().min( 1 ),
+    tools: z.record( z.string(), z.record( z.string(), z.unknown() ) ).optional()
+  } ).loose()
+} ).loose();
+
 const toolConfigSchema = z.record( z.string(), z.unknown() );
 
+const parseModelPrompt = prompt => {
+  const result = modelPromptSchema.safeParse( prompt );
+  if ( !result.success ) {
+    throw new ValidationError( `Invalid model prompt config: ${z.prettifyError( result.error )}` );
+  }
+  return result.data.config;
+};
+
+const parseToolsPrompt = prompt => {
+  const result = toolsPromptSchema.safeParse( prompt );
+  if ( !result.success ) {
+    throw new ValidationError( `Invalid tools prompt config: ${z.prettifyError( result.error )}` );
+  }
+  return result.data.config;
+};
+
+/**
+ * Register or override an AI SDK provider factory by name.
+ *
+ * @param {string} name - Provider name used in prompt frontmatter
+ * @param {Function} providerFn - Factory function that receives a model id
+ * @returns {void}
+ */
 export function registerProvider( name, providerFn ) {
   const result = registerProviderSchema.safeParse( { name, providerFn } );
   if ( !result.success ) {
@@ -41,25 +78,23 @@ export function registerProvider( name, providerFn ) {
   providers[name] = providerFn;
 }
 
+/**
+ * List all currently registered provider names.
+ *
+ * @returns {string[]} Provider names
+ */
 export const getRegisteredProviders = () => Object.keys( providers );
 
-export function loadModel( prompt ) {
-  const config = prompt?.config;
-
-  if ( !config ) {
-    throw new Error( 'Prompt is missing config object' );
-  }
-
+/**
+ * Load a text model from a loaded prompt config.
+ *
+ * @param {object} prompt - Loaded prompt object with `config.provider` and `config.model`
+ * @returns {unknown} AI SDK language model
+ */
+export function loadTextModel( prompt ) {
+  const config = parseModelPrompt( prompt );
   const { provider: providerName, model: modelName } = config;
 
-  if ( !providerName ) {
-    throw new Error( 'Prompt config is missing "provider" field' );
-  }
-
-  if ( !modelName ) {
-    throw new Error( 'Prompt config is missing "model" field' );
-  }
-
   const provider = providers[providerName];
 
   if ( !provider ) {
@@ -70,34 +105,42 @@ export function loadModel( prompt ) {
   return provider( modelName );
 }
 
-export function loadTools( prompt ) {
-  const config = prompt?.config;
-
-  if ( !config ) {
-    throw new Error( 'Prompt is missing config object' );
-  }
+/**
+ * Load an image model from a loaded prompt config.
+ *
+ * @param {object} prompt - Loaded prompt object with `config.provider` and `config.model`
+ * @returns {unknown} AI SDK image model
+ */
+export function loadImageModel( prompt ) {
+  const config = parseModelPrompt( prompt );
+  const { provider: providerName, model: modelName } = config;
 
-  const { tools: toolsConfig, provider: providerName } = config;
+  const provider = providers[providerName];
 
-  if ( !toolsConfig ) {
-    return null;
+  if ( !provider ) {
+    const availableProviders = Object.keys( providers ).join( ', ' );
+    throw new Error( `Invalid provider "${providerName}". Valid providers: ${availableProviders}` );
   }
 
-  if ( Array.isArray( toolsConfig ) ) {
-    throw new Error(
-      'tools must be an object with tool configurations, got array. ' +
-      'Use "tools: { googleSearch: {} }" not "tools: [googleSearch]"'
-    );
+  const imageModelFactory = provider.image ?? provider.imageModel;
+  if ( typeof imageModelFactory !== 'function' ) {
+    throw new Error( `Provider "${providerName}" does not support image models.` );
   }
 
-  if ( typeof toolsConfig !== 'object' ) {
-    throw new Error(
-      `tools must be an object, got ${typeof toolsConfig}. ` +
-      'Use "tools: { googleSearch: {} }"'
-    );
-  }
+  return imageModelFactory( modelName );
+}
+
+/**
+ * Load provider-specific tools configured in a prompt.
+ *
+ * @param {object} prompt - Loaded prompt object with `config.provider` and optional `config.tools`
+ * @returns {Record<string, unknown> | null} AI SDK tools, or null when none are configured
+ */
+export function loadTools( prompt ) {
+  const config = parseToolsPrompt( prompt );
+  const { tools: toolsConfig, provider: providerName } = config;
 
-  if ( Object.keys( toolsConfig ).length === 0 ) {
+  if ( !toolsConfig || Object.keys( toolsConfig ).length === 0 ) {
     return null;
   }
 

package/src/ai_model.spec.js

@@ -2,6 +2,7 @@ import { it, expect, vi, afterEach, describe } from 'vitest';
 
 const providerFactoryOptions = vi.hoisted( () => ( {} ) );
 const openaiImpl = vi.hoisted( () => vi.fn( model => `openai:${model}` ) );
+const openaiImageImpl = vi.hoisted( () => vi.fn( model => `openai-image:${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}` ) );
@@ -11,6 +12,7 @@ const vertexImpl = vi.hoisted( () => vi.fn( model => `vertex:${model}` ) );
 // OpenAI mock with tools support
 vi.mock( '@ai-sdk/openai', () => {
   const openaiMock = ( ...values ) => openaiImpl( ...values );
+  openaiMock.image = ( ...values ) => openaiImageImpl( ...values );
   openaiMock.tools = {
     webSearch: ( config = {} ) => ( { type: 'webSearch', config } )
   };
@@ -93,14 +95,14 @@ vi.mock( '@ai-sdk/google-vertex', () => {
   };
 } );
 
-import { loadModel, loadTools, registerProvider, getRegisteredProviders, providers, builtInProviders } from './ai_model.js';
+import { loadImageModel, loadTextModel, loadTools, registerProvider, getRegisteredProviders, providers, builtInProviders } from './ai_model.js';
 
 afterEach( async () => {
   await vi.resetModules();
   vi.clearAllMocks();
 } );
 
-describe( 'loadModel', () => {
+describe( 'loadTextModel', () => {
   it( 'initializes built-in providers with custom fetch', () => {
     expect( providerFactoryOptions ).toMatchObject( {
       azure: { fetch: expect.any( Function ) },
@@ -113,7 +115,7 @@ describe( 'loadModel', () => {
   } );
 
   it( 'loads model using selected provider', () => {
-    const result = loadModel( { config: { provider: 'openai', model: 'gpt-4o-mini' } } );
+    const result = loadTextModel( { config: { provider: 'openai', model: 'gpt-4o-mini' } } );
 
     expect( result ).toBe( 'openai:gpt-4o-mini' );
     expect( openaiImpl ).toHaveBeenCalledWith( 'gpt-4o-mini' );
@@ -122,20 +124,36 @@ describe( 'loadModel', () => {
   } );
 
   it( 'loads model using bedrock provider', () => {
-    const result = loadModel( { config: { provider: 'bedrock', model: 'anthropic.claude-sonnet-4-20250514-v1:0' } } );
+    const result = loadTextModel( { config: { provider: 'bedrock', model: 'anthropic.claude-sonnet-4-20250514-v1:0' } } );
 
     expect( result ).toBe( 'bedrock:anthropic.claude-sonnet-4-20250514-v1:0' );
     expect( bedrockImpl ).toHaveBeenCalledWith( 'anthropic.claude-sonnet-4-20250514-v1:0' );
   } );
 
   it( 'loads model using perplexity provider', () => {
-    const result = loadModel( { config: { provider: 'perplexity', model: 'sonar-pro' } } );
+    const result = loadTextModel( { config: { provider: 'perplexity', model: 'sonar-pro' } } );
 
     expect( result ).toBe( 'perplexity:sonar-pro' );
     expect( perplexityImpl ).toHaveBeenCalledWith( 'sonar-pro' );
   } );
 } );
 
+describe( 'loadImageModel', () => {
+  it( 'loads image model using provider image factory', () => {
+    const result = loadImageModel( { config: { provider: 'openai', model: 'gpt-image-1' } } );
+
+    expect( result ).toBe( 'openai-image:gpt-image-1' );
+    expect( openaiImageImpl ).toHaveBeenCalledWith( 'gpt-image-1' );
+    expect( openaiImpl ).not.toHaveBeenCalled();
+  } );
+
+  it( 'throws a clear error when provider does not support image models', () => {
+    expect( () => loadImageModel( {
+      config: { provider: 'azure', model: 'gpt-image-1' }
+    } ) ).toThrow( 'Provider "azure" does not support image models.' );
+  } );
+} );
+
 describe( 'loadTools', () => {
   // Category 1: Basic Functionality (5 tests)
   describe( 'Basic Functionality', () => {
@@ -593,7 +611,7 @@ describe( 'loadTools', () => {
           provider: 'vertex',
           tools: [ 'googleSearch', 'urlContext' ]
         }
-      } ) ).toThrow( 'tools must be an object with tool configurations, got array' );
+      } ) ).toThrow( /Invalid tools prompt config/ );
     } );
 
     it( 'throws error for string format', () => {
@@ -602,7 +620,7 @@ describe( 'loadTools', () => {
           provider: 'vertex',
           tools: 'googleSearch'
         }
-      } ) ).toThrow( 'tools must be an object' );
+      } ) ).toThrow( /Invalid tools prompt config/ );
     } );
 
     it( 'throws error for number format', () => {
@@ -611,7 +629,7 @@ describe( 'loadTools', () => {
           provider: 'vertex',
           tools: 123
         }
-      } ) ).toThrow( 'tools must be an object' );
+      } ) ).toThrow( /Invalid tools prompt config/ );
     } );
 
     it( 'throws error for provider without tools support', () => {
@@ -656,7 +674,7 @@ describe( 'loadTools', () => {
           provider: 'vertex',
           tools: { googleSearch: null }
         }
-      } ) ).toThrow( /Invalid config for tool "googleSearch".*expected record, received null/s );
+      } ) ).toThrow( /Invalid tools prompt config.*expected record, received null/s );
     } );
 
     it( 'throws error when tool config is a string', () => {
@@ -665,7 +683,7 @@ describe( 'loadTools', () => {
           provider: 'vertex',
           tools: { googleSearch: 'MODE_DYNAMIC' }
         }
-      } ) ).toThrow( /Invalid config for tool "googleSearch".*expected record, received string/s );
+      } ) ).toThrow( /Invalid tools prompt config.*expected record, received string/s );
     } );
 
     it( 'throws error for unknown tool on Bedrock with dynamic tool listing', () => {
@@ -753,11 +771,11 @@ describe( 'registerProvider', () => {
     Object.assign( providers, builtInProviders );
   } );
 
-  it( 'registers a custom provider and uses it in loadModel', () => {
+  it( 'registers a custom provider and uses it in loadTextModel', () => {
     const customProvider = vi.fn( model => `custom:${model}` );
     registerProvider( 'custom', customProvider );
 
-    const result = loadModel( { config: { provider: 'custom', model: 'my-model' } } );
+    const result = loadTextModel( { config: { provider: 'custom', model: 'my-model' } } );
 
     expect( result ).toBe( 'custom:my-model' );
     expect( customProvider ).toHaveBeenCalledWith( 'my-model' );
@@ -767,7 +785,7 @@ describe( 'registerProvider', () => {
     const overrideOpenai = vi.fn( model => `override:${model}` );
     registerProvider( 'openai', overrideOpenai );
 
-    const result = loadModel( { config: { provider: 'openai', model: 'gpt-custom' } } );
+    const result = loadTextModel( { config: { provider: 'openai', model: 'gpt-custom' } } );
 
     expect( result ).toBe( 'override:gpt-custom' );
   } );

package/src/ai_sdk.js

@@ -1,113 +1,89 @@
-import { loadModel, loadTools } from './ai_model.js';
+import { types as utilTypes } from 'node:util';
 import * as AI from 'ai';
 import { stepCountIs } from 'ai';
-import { validateGenerateTextArgs, validateStreamTextArgs } from './validations.js';
-import { loadPrompt } from './prompt_loader.js';
-import { buildSystemSkillsVar, buildLoadSkillTool, loadPromptSkills, loadColocatedSkills } from './skill.js';
+import { ValidationError } from '@outputai/core';
+import { validateGenerateTextArgs, validateStreamTextArgs, validateGenerateImageArgs } from './validations.js';
+import { loadPrompt } from './prompt/loader.js';
 import { startTrace, endTraceWithError } from './utils/trace.js';
-import { wrapTextResponse, wrapStreamOnFinishResponse } from './utils/response_wrappers.js';
-
-export const loadAiSdkOptionsFromPrompt = prompt => {
-  const options = {
-    model: loadModel( prompt ),
-    messages: prompt.messages,
-    providerOptions: prompt.config.providerOptions
-  };
-
-  if ( Number.isFinite( prompt.config.temperature ) ) {
-    options.temperature = prompt.config.temperature;
-  }
-
-  if ( prompt.config.maxTokens ) {
-    options.maxOutputTokens = prompt.config.maxTokens;
-  }
-
-  const tools = loadTools( prompt );
-  if ( tools ) {
-    options.tools = tools;
-  }
-
-  return options;
-};
-
-export const hydratePromptTemplate = ( prompt, variables, promptDir, callerSkills, callerTools = {} ) => {
-  const meta = loadPrompt( prompt, variables, promptDir );
-
-  // Resolve skills: explicit frontmatter paths > colocated auto-discovery
-  const hasExplicitSkills = meta.config.skills && meta.promptFileDir;
-  const frontmatterSkills = hasExplicitSkills ?
-    loadPromptSkills( meta.config.skills, meta.promptFileDir ) :
-    [];
-  const autoSkills = !hasExplicitSkills && meta.promptFileDir ?
-    loadColocatedSkills( meta.promptFileDir ) :
-    [];
-  const resolvedSkills = [ ...frontmatterSkills, ...autoSkills, ...callerSkills ];
-
-  const tools = resolvedSkills.length > 0 ?
-    { load_skill: buildLoadSkillTool( resolvedSkills ), ...callerTools } :
-    callerTools;
-
-  const skillsMessage = resolvedSkills.length > 0 ?
-    { role: 'system', content: buildSystemSkillsVar( resolvedSkills ) } :
-    null;
-
-  if ( skillsMessage ) {
-    // Merge into existing system message to avoid provider errors with multiple system messages
-    const systemMsg = meta.messages.find( m => m.role === 'system' );
-    if ( systemMsg ) {
-      systemMsg.content = `${systemMsg.content}\n\n${skillsMessage.content}`;
-    } else {
-      meta.messages.unshift( skillsMessage );
-    }
-  }
-
-  return { loadedPrompt: meta, allVariables: variables, tools };
-};
+import { wrapTextResponse, wrapStreamOnFinishResponse, wrapImageResponse } from './utils/response_wrappers.js';
+import { loadAiSdkTextOptions, loadAiSdkImageOptions } from './ai_sdk_options.js';
+import { prepareTextPrompt } from './prompt/prepare_text.js';
+import { mapAiError } from './utils/error_handler.js';
 
 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, aiSdkArgs.tools );
-  const hasTools = Object.keys( tools ).length > 0;
+  validateGenerateTextArgs( { prompt, variables, promptDir, skills, maxSteps } );
 
-  validateGenerateTextArgs( { prompt, variables: allVariables } );
+  const parsedSkills = typeof skills === 'function' ? await skills( variables ) : skills;
+  const { loadedPrompt, tools } = prepareTextPrompt( { prompt, variables, promptDir, skills: parsedSkills, tools: aiSdkArgs.tools } );
 
-  const traceId = startTrace( { name: 'generateText', prompt, variables: allVariables, loadedPrompt } );
+  const traceId = startTrace( { name: 'generateText', prompt, variables, loadedPrompt } );
   const { model: modelId } = loadedPrompt.config;
 
   try {
     const response = await AI.generateText( {
-      ...loadAiSdkOptionsFromPrompt( loadedPrompt ),
+      ...loadAiSdkTextOptions( loadedPrompt ),
       maxRetries: 0,
       ...aiSdkArgs,
-      ...( hasTools ? { tools } : {} ),
-      ...( hasTools && !aiSdkArgs.stopWhen ? { stopWhen: stepCountIs( maxSteps ) } : {} )
+      ...( tools && { tools } ),
+      ...( tools && !aiSdkArgs.stopWhen ? { stopWhen: stepCountIs( maxSteps ) } : {} )
     } );
     return wrapTextResponse( { traceId, modelId, response } );
-  } catch ( error ) {
+  } catch ( originalError ) {
+    const error = mapAiError( originalError );
     endTraceWithError( { traceId, error } );
     throw error;
   }
 }
 
-export function streamText( { prompt, variables, onFinish, onError, ...aiSdkArgs } ) {
-  validateStreamTextArgs( { prompt, variables } );
-  const loadedPrompt = loadPrompt( prompt, variables );
+export function streamText( { prompt, variables, promptDir, skills = [], maxSteps = 10, onFinish, onError: _onError, ...aiSdkArgs } ) {
+  validateStreamTextArgs( { prompt, variables, promptDir, skills, maxSteps } );
+
+  const parsedSkills = typeof skills === 'function' ? skills( variables ) : skills;
+  if ( utilTypes.isPromise( parsedSkills ) ) {
+    throw new ValidationError( 'streamText() skills must be synchronous because streamText() returns a stream immediately.' );
+  }
+  const { loadedPrompt, tools } = prepareTextPrompt( { prompt, variables, promptDir, skills: parsedSkills, tools: aiSdkArgs.tools } );
+
   const traceId = startTrace( { name: 'streamText', prompt, variables, loadedPrompt } );
   const { model: modelId } = loadedPrompt.config;
 
   try {
     return AI.streamText( {
-      ...loadAiSdkOptionsFromPrompt( loadedPrompt ),
+      ...loadAiSdkTextOptions( loadedPrompt ),
       maxRetries: 0,
       ...aiSdkArgs,
+      ...( tools && { tools } ),
+      ...( tools && !aiSdkArgs.stopWhen ? { stopWhen: stepCountIs( maxSteps ) } : {} ),
       ...wrapStreamOnFinishResponse( { traceId, modelId, onFinish } ),
       onError( event ) {
-        endTraceWithError( { traceId, error: event.error } );
-        onError?.( event );
+        const error = mapAiError( event.error );
+        endTraceWithError( { traceId, error } );
+        _onError?.( { ...event, error } );
       }
     } );
-  } catch ( error ) {
+  } catch ( originalError ) {
+    const error = mapAiError( originalError );
+    endTraceWithError( { traceId, error } );
+    throw error;
+  }
+}
+
+export async function generateImage( { prompt, variables, promptDir, images, mask, ...aiSdkArgs } ) {
+  validateGenerateImageArgs( { prompt, variables, promptDir, images, mask } );
+
+  const loadedPrompt = loadPrompt( prompt, variables, promptDir );
+  const traceId = startTrace( { name: 'generateImage', prompt, variables, loadedPrompt } );
+  const { model: modelId } = loadedPrompt.config;
+
+  try {
+    const response = await AI.generateImage( {
+      ...loadAiSdkImageOptions( { prompt: loadedPrompt, images, mask } ),
+      maxRetries: 0,
+      ...aiSdkArgs
+    } );
+    return wrapImageResponse( { traceId, modelId, response } );
+  } catch ( originalError ) {
+    const error = mapAiError( originalError );
     endTraceWithError( { traceId, error } );
     throw error;
   }