@output.ai/llm 0.0.15 → 0.1.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@output.ai/llm",
-  "version": "0.0.15",
+  "version": "0.1.0",
   "description": "Framework abstraction to interact with LLM models",
   "type": "module",
   "main": "src/index.js",

package/src/ai_model.js

@@ -1,29 +1,9 @@
 import { anthropic } from '@ai-sdk/anthropic';
 import { azure } from '@ai-sdk/azure';
 import { openai } from '@ai-sdk/openai';
-import { ValidationError, z } from '@output.ai/core';
 
 const providers = { azure, anthropic, openai };
 
-const promptSchema = z.object( {
-  config: z.object( {
-    provider: z.enum( [ 'anthropic', 'azure', 'openai' ] ),
-    model: z.string(),
-    temperature: z.number().optional(),
-    max_tokens: z.number().optional()
-  } ),
-  messages: z.array(
-    z.object( {
-      role: z.string(),
-      content: z.string()
-    } )
-  )
-} );
-
 export const loadModel = prompt => {
-  const result = promptSchema.safeParse( prompt );
-  if ( !result.success ) {
-    throw new ValidationError( `Invalid prompt object: ${result.error.message}` );
-  }
   return providers[prompt.config.provider]( prompt.config.model );
 };

package/src/ai_model.spec.js

@@ -0,0 +1,33 @@
+import { it, expect, vi, afterEach } from 'vitest';
+
+const openaiImpl = vi.fn( model => `openai:${model}` );
+const azureImpl = vi.fn( model => `azure:${model}` );
+const anthropicImpl = vi.fn( model => `anthropic:${model}` );
+
+vi.mock( '@ai-sdk/openai', () => ( {
+  openai: ( ...values ) => openaiImpl( ...values )
+} ) );
+
+vi.mock( '@ai-sdk/azure', () => ( {
+  azure: ( ...values ) => azureImpl( ...values )
+} ) );
+
+vi.mock( '@ai-sdk/anthropic', () => ( {
+  anthropic: ( ...values ) => anthropicImpl( ...values )
+} ) );
+
+import { loadModel } from './ai_model.js';
+
+afterEach( async () => {
+  await vi.resetModules();
+  vi.clearAllMocks();
+} );
+
+it( 'loads model using selected provider', () => {
+  const result = loadModel( { config: { provider: 'openai', model: 'gpt-4o-mini' } } );
+
+  expect( result ).toBe( 'openai:gpt-4o-mini' );
+  expect( openaiImpl ).toHaveBeenCalledWith( 'gpt-4o-mini' );
+  expect( azureImpl ).not.toHaveBeenCalled();
+  expect( anthropicImpl ).not.toHaveBeenCalled();
+} );

package/src/ai_sdk.js

@@ -1,8 +1,9 @@
 import { Tracing } from '@output.ai/core/tracing';
 import { loadModel } from './ai_model.js';
 import * as AI from 'ai';
+import { validateGenerateTextArgs, validateGenerateObjectArgs, validateGenerateArrayArgs, validateGenerateEnumArgs } from './validations.js';
 
-const generationWrapper = async ( traceId, fn ) => {
+const traceWrapper = async ( traceId, fn ) => {
   try {
     const result = await fn();
     Tracing.addEventEnd( { id: traceId, details: result } );
@@ -13,33 +14,103 @@ const generationWrapper = async ( traceId, fn ) => {
   }
 };
 
-export async function generateText( { prompt, ...nativeAiSdkArgs } ) {
+const extraAiSdkOptionsFromPrompt = prompt => ( {
+  model: loadModel( prompt ),
+  messages: prompt.messages,
+  ...( prompt.config.temperature && { temperature: prompt.config.temperature } ),
+  ...( prompt.config.max_tokens && { maxOutputTokens: prompt.config.maxTokens } ),
+  providerOptions: prompt.providerOptions
+} );
+
+/**
+ * Use an LLM model to generate text.
+ *
+ * @param {object} args - Generation arguments
+ * @param {Prompt} args.prompt - `@output.ai/prompt` Prompt object
+ * @returns {Promise<string>} Generated text
+ */
+export async function generateText( { prompt } ) {
+  validateGenerateTextArgs( { prompt } );
   const traceId = `generateText-${Date.now()}`;
-  Tracing.addEventStart( { kind: 'llm', name: 'generateText', id: traceId, details: { prompt, nativeAiSdkArgs } } );
-
-  return generationWrapper( traceId, async () => {
-    return ( await AI.generateText( {
-      model: loadModel( prompt ),
-      messages: prompt.messages,
-      temperature: prompt.config.temperature,
-      maxOutputTokens: prompt.config.max_tokens ?? 64000,
-      ...nativeAiSdkArgs
-    } ) ).text;
-  } );
+  Tracing.addEventStart( { kind: 'llm', name: 'generateText', id: traceId, details: { prompt } } );
+
+  return traceWrapper( traceId, async () =>
+    AI.generateText( extraAiSdkOptionsFromPrompt( prompt ) ).then( r => r.text )
+  );
 }
 
-export async function generateObject( { prompt, ...nativeAiSdkArgs } ) {
+/**
+ * Use an LLM model to generate an object with a fixed schema.
+ *
+ * @param {object} args - Generation arguments
+ * @param {Prompt} args.prompt - `@output.ai/prompt` Prompt object
+ * @param {z.ZodObject} args.schema - Output schema
+ * @param {string} [args.schemaName] - Output schema name
+ * @param {string} [args.schemaDescription] - Output schema description
+ * @returns {Promise<object>} Object matching the provided schema
+ */
+export async function generateObject( args ) {
+  validateGenerateObjectArgs( args );
+  const { prompt, schema, schemaName, schemaDescription } = args;
+
   const traceId = `generateObject-${Date.now()}`;
-  Tracing.addEventStart( { kind: 'llm', name: 'generateObject', id: traceId, details: { prompt, nativeAiSdkArgs } } );
-
-  return generationWrapper( traceId, async () => {
-    return ( await AI.generateObject( {
-      model: loadModel( prompt ),
-      output: nativeAiSdkArgs.object ?? 'object',
-      messages: prompt.messages,
-      temperature: prompt.config.temperature,
-      ...( prompt.config.max_tokens && { maxOutputTokens: prompt.config.max_tokens } ),
-      ...nativeAiSdkArgs
-    } ) ).object;
-  } );
+  Tracing.addEventStart( { kind: 'llm', name: 'generateObject', id: traceId, details: args } );
+
+  return traceWrapper( traceId, async () =>
+    AI.generateObject( {
+      output: 'object',
+      schema,
+      schemaName,
+      schemaDescription,
+      ...extraAiSdkOptionsFromPrompt( prompt )
+    } ).then( r => r.object )
+  );
+}
+
+/**
+ * Use an LLM model to generate an array of values with a fixed schema.
+ *
+ * @param {object} args - Generation arguments
+ * @param {Prompt} args.prompt - `@output.ai/prompt` Prompt object
+ * @param {z.ZodType} args.schema - Output schema (array item)
+ * @param {string} [args.schemaName] - Output schema name
+ * @param {string} [args.schemaDescription] - Output schema description
+ * @returns {Promise<object>} Array where each element matches the schema
+ */
+export async function generateArray( args ) {
+  validateGenerateArrayArgs( args );
+  const { prompt, schema, schemaName, schemaDescription } = args;
+
+  const traceId = `generateArray-${Date.now()}`;
+  Tracing.addEventStart( { kind: 'llm', name: 'generateArray', id: traceId, details: args } );
+
+  return traceWrapper( traceId, async () =>
+    AI.generateObject( {
+      output: 'array',
+      schema,
+      schemaName,
+      schemaDescription,
+      ...extraAiSdkOptionsFromPrompt( prompt )
+    } ).then( r => r.object )
+  );
+}
+
+/**
+ * Use an LLM model to generate a result from an enum (array of string values).
+ *
+ * @param {object} args - Generation arguments
+ * @param {Prompt} args.prompt - `@output.ai/prompt` Prompt object
+ * @param {string[]} args.enum - Allowed values for the generation
+ * @returns {Promise<string>} One of the provided enum values
+ */
+export async function generateEnum( args ) {
+  validateGenerateEnumArgs( args );
+  const { prompt, enum: _enum } = args;
+
+  const traceId = `generateEnum-${Date.now()}`;
+  Tracing.addEventStart( { kind: 'llm', name: 'generateEnum', id: traceId, details: args } );
+
+  return traceWrapper( traceId, async () =>
+    AI.generateObject( { output: 'enum', enum: _enum, ...extraAiSdkOptionsFromPrompt( prompt ) } ).then( r => r.object )
+  );
 }

package/src/ai_sdk.spec.js

@@ -0,0 +1,156 @@
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { z } from '@output.ai/core';
+
+const tracingSpies = {
+  addEventStart: vi.fn(),
+  addEventEnd: vi.fn(),
+  addEventError: vi.fn()
+};
+vi.mock( '@output.ai/core/tracing', () => ( { Tracing: tracingSpies } ), { virtual: true } );
+
+const loadModelImpl = vi.fn();
+vi.mock( './ai_model.js', () => ( {
+  loadModel: ( ...values ) => loadModelImpl( ...values )
+} ) );
+
+const aiFns = {
+  generateText: vi.fn(),
+  generateObject: vi.fn()
+};
+vi.mock( 'ai', () => ( aiFns ) );
+
+const validators = {
+  validateGenerateTextArgs: vi.fn(),
+  validateGenerateObjectArgs: vi.fn(),
+  validateGenerateArrayArgs: vi.fn(),
+  validateGenerateEnumArgs: vi.fn()
+};
+vi.mock( './validations.js', () => ( validators ) );
+
+const importSut = async () => import( './ai_sdk.js' );
+
+const basePrompt = {
+  config: { provider: 'openai', model: 'gpt-4o-mini', temperature: 0.3 },
+  messages: [ { role: 'user', content: 'Hi' } ],
+  providerOptions: { thinking: { enabled: true } }
+};
+
+beforeEach( () => {
+  tracingSpies.addEventStart.mockClear();
+  tracingSpies.addEventEnd.mockClear();
+  tracingSpies.addEventError.mockClear();
+
+  loadModelImpl.mockReset().mockReturnValue( 'MODEL' );
+
+  aiFns.generateText.mockReset().mockResolvedValue( { text: 'TEXT' } );
+  aiFns.generateObject.mockReset().mockResolvedValue( { object: 'OBJECT' } );
+
+  validators.validateGenerateTextArgs.mockClear();
+  validators.validateGenerateObjectArgs.mockClear();
+  validators.validateGenerateArrayArgs.mockClear();
+  validators.validateGenerateEnumArgs.mockClear();
+} );
+
+afterEach( async () => {
+  await vi.resetModules();
+  vi.clearAllMocks();
+} );
+
+describe( 'ai_sdk', () => {
+  it( 'generateText: validates, traces, calls AI and returns text', async () => {
+    const { generateText } = await importSut();
+    const result = await generateText( { prompt: basePrompt } );
+
+    expect( validators.validateGenerateTextArgs ).toHaveBeenCalledWith( { prompt: basePrompt } );
+    expect( tracingSpies.addEventStart ).toHaveBeenCalledTimes( 1 );
+    expect( tracingSpies.addEventEnd ).toHaveBeenCalledTimes( 1 );
+
+    expect( loadModelImpl ).toHaveBeenCalledWith( basePrompt );
+    expect( aiFns.generateText ).toHaveBeenCalledWith( {
+      model: 'MODEL',
+      messages: basePrompt.messages,
+      temperature: 0.3,
+      providerOptions: basePrompt.providerOptions
+    } );
+    expect( result ).toBe( 'TEXT' );
+  } );
+
+  it( 'generateObject: validates, traces, calls AI with output object and returns object', async () => {
+    const { generateObject } = await importSut();
+    aiFns.generateObject.mockResolvedValueOnce( { object: { a: 1 } } );
+
+    const schema = z.object( { a: z.number() } );
+    const result = await generateObject( {
+      prompt: basePrompt,
+      schema,
+      schemaName: 'Thing',
+      schemaDescription: 'A thing'
+    } );
+
+    expect( validators.validateGenerateObjectArgs ).toHaveBeenCalled();
+    expect( tracingSpies.addEventStart ).toHaveBeenCalledTimes( 1 );
+    expect( tracingSpies.addEventEnd ).toHaveBeenCalledTimes( 1 );
+
+    expect( aiFns.generateObject ).toHaveBeenCalledWith( {
+      output: 'object',
+      schema,
+      schemaName: 'Thing',
+      schemaDescription: 'A thing',
+      model: 'MODEL',
+      messages: basePrompt.messages,
+      temperature: 0.3,
+      providerOptions: basePrompt.providerOptions
+    } );
+    expect( result ).toEqual( { a: 1 } );
+  } );
+
+  it( 'generateArray: validates, traces, calls AI (item schema) and returns array', async () => {
+    const { generateArray } = await importSut();
+    aiFns.generateObject.mockResolvedValueOnce( { object: [ 1, 2 ] } );
+
+    const schema = z.number();
+    const result = await generateArray( {
+      prompt: basePrompt,
+      schema,
+      schemaName: 'Numbers',
+      schemaDescription: 'Two numbers'
+    } );
+
+    expect( validators.validateGenerateArrayArgs ).toHaveBeenCalled();
+    expect( tracingSpies.addEventStart ).toHaveBeenCalledTimes( 1 );
+    expect( tracingSpies.addEventEnd ).toHaveBeenCalledTimes( 1 );
+
+    expect( aiFns.generateObject ).toHaveBeenCalledWith( {
+      output: 'array',
+      schema,
+      schemaName: 'Numbers',
+      schemaDescription: 'Two numbers',
+      model: 'MODEL',
+      messages: basePrompt.messages,
+      temperature: 0.3,
+      providerOptions: basePrompt.providerOptions
+    } );
+    expect( result ).toEqual( [ 1, 2 ] );
+  } );
+
+  it( 'generateEnum: validates, traces, calls AI with output enum and returns value', async () => {
+    const { generateEnum } = await importSut();
+    aiFns.generateObject.mockResolvedValueOnce( { object: 'B' } );
+
+    const result = await generateEnum( { prompt: basePrompt, enum: [ 'A', 'B', 'C' ] } );
+
+    expect( validators.validateGenerateEnumArgs ).toHaveBeenCalled();
+    expect( tracingSpies.addEventStart ).toHaveBeenCalledTimes( 1 );
+    expect( tracingSpies.addEventEnd ).toHaveBeenCalledTimes( 1 );
+
+    expect( aiFns.generateObject ).toHaveBeenCalledWith( {
+      output: 'enum',
+      enum: [ 'A', 'B', 'C' ],
+      model: 'MODEL',
+      messages: basePrompt.messages,
+      temperature: 0.3,
+      providerOptions: basePrompt.providerOptions
+    } );
+    expect( result ).toBe( 'B' );
+  } );
+} );

package/src/index.d.ts

@@ -1,110 +1,82 @@
-import type * as AiTypes from 'ai';
-import type { z as CoreZ } from '@output.ai/core';
+import type { z } from '@output.ai/core';
 import type { Prompt } from '@output.ai/prompt';
 
 export type { Prompt };
 
-type NativeGenerateTextArgs = Parameters<typeof AiTypes.generateText>[0];
-type NativeGenerateObjectArgs = Parameters<typeof AiTypes.generateObject>[0];
-
-/**
- * Simplify types into a plain object while preserving unions
- * (distributes over union members instead of collapsing their keys)
- */
-type Simplify<T> = T extends unknown ? { [K in keyof T]: T[K] } & {} : never;
-
-/**
- * Replace keys K in T with V, preserving unions by distributing over T
- */
-type Replace<T, K extends PropertyKey, V> = T extends unknown
-  ? Omit<T, Extract<K, keyof T>> & { [P in K]: V }
-  : never;
-
-/**
- * Text generation arguments
- * Include all native AI SDK generateText options and a Prompt object from `@output.ai/prompt`
- */
-export type GenerateTextArgs = Simplify<
-  Replace<Partial<NativeGenerateTextArgs>, 'prompt', Prompt>
->;
-
-/**
- * Allow schemas from @output.ai/core's zod in addition to AI SDK's accepted schema types.
- */
-type WithCoreZodSchema<T> = T extends unknown
-  ? T extends { schema?: infer S }
-    ? Omit<T, 'schema'> & { schema?: S | CoreZ.ZodTypeAny }
-    : T
-  : never;
-
-/**
- * Object generation arguments
- * Include all native AI SDK generateObject options and a Prompt object from `@output.ai/prompt`
- */
-export type GenerateObjectArgs = Simplify<
-  Replace<Partial<WithCoreZodSchema<NativeGenerateObjectArgs>>, 'prompt', Prompt>
->;
-
 /**
- * Use a LLM Model to generate text
- *
- * This function a wrapper over AI SDK's generateText function.
- *
- * It accepts the same arguments of the the original function, plus a "Prompt" object, generated using the `output.ai/prompt`.
+ * Use an LLM model to generate text.
  *
- * The Prompt object will set `model`, `messages`, `temperature` and `max_tokens`, however all these can be overwritten by their AI SDK native argument values.
+ * This function is a wrapper over the AI SDK's `generateText`.
+ * The `Prompt` sets `model`, `messages`, `temperature`, `max_tokens`, and `provider_options`.
  *
- * @param {GenerateTextArgs} args - Generation arguments
- * @returns {Promise<string>}
+ * @param {object} args - Generation arguments
+ * @param {Prompt} args.prompt - `@output.ai/prompt` Prompt object
+ * @returns {Promise<string>} Generated text
  */
 export function generateText(
-  args: GenerateTextArgs
+  args: {
+    prompt: Prompt
+  }
 ): Promise<string>;
 
 /**
- * Use a LLM Model to generate object
+ * Use an LLM model to generate an object with a fixed schema.
  *
- * This function a wrapper over AI SDK's generateObject function.
+ * This function is a wrapper over the AI SDK's `generateObject`.
+ * The `Prompt` sets `model`, `messages`, `temperature`, `max_tokens`, and `provider_options`.
  *
- * It accepts the same arguments of the the original function, plus a "Prompt" object, generated using the `output.ai/prompt`.
- *
- * The Prompt object will set `model`, `messages`, `temperature` and `max_tokens`, however all these can be overwritten by their AI SDK native argument values.
- *
- * @param {GenerateObjectArgs} args - Generation arguments
- * @returns {Promise<object>} An object matching the provided schema
+ * @param {object} args - Generation arguments
+ * @param {Prompt} args.prompt - `@output.ai/prompt` Prompt object
+ * @param {z.ZodObject} args.schema - Output schema
+ * @param {string} [args.schemaName] - Output schema name
+ * @param {string} [args.schemaDescription] - Output schema description
+ * @returns {Promise<object>} Object matching the provided schema
  */
-export function generateObject<A extends GenerateObjectArgs & { schema: CoreZ.ZodTypeAny }>(
-  args: A
-): Promise<CoreZ.infer<A['schema']>>;
+export function generateObject<TSchema extends z.ZodObject>(
+  args: {
+    prompt: Prompt,
+    schema?: TSchema,
+    schemaName?: string,
+    schemaDescription?: string
+  }
+): Promise<z.infer<TSchema>>;
 
 /**
- * Use a LLM Model to generate object
- *
- * This function a wrapper over AI SDK's generateObject function.
+ * Use an LLM model to generate an array of values with a fixed schema.
  *
- * It accepts the same arguments of the the original function, plus a "Prompt" object, generated using the `output.ai/prompt`.
+ * This function is a wrapper over the AI SDK's `generateObject` with `output: 'array'`.
+ * The `Prompt` sets `model`, `messages`, `temperature`, `max_tokens`, and `provider_options`.
  *
- * The Prompt object will set `model`, `messages`, `temperature` and `max_tokens`, however all these can be overwritten by their AI SDK native argument values.
- *
- * @param {GenerateObjectArgs} args - Generation arguments
- * @returns {Promise<object>} An object matching the provided enum
+ * @param {object} args - Generation arguments
+ * @param {Prompt} args.prompt - `@output.ai/prompt` Prompt object
+ * @param {z.ZodType} args.schema - Output schema (array item)
+ * @param {string} [args.schemaName] - Output schema name
+ * @param {string} [args.schemaDescription] - Output schema description
+ * @returns {Promise<object>} Array where each element matches the schema
  */
-export function generateObject<A extends GenerateObjectArgs & { enum: readonly unknown[]; output: 'enum' }>(
-  args: A
-): Promise<A['enum'][number]>;
+export function generateArray<TSchema extends z.ZodType>(
+  args: {
+    prompt: Prompt,
+    schema?: TSchema,
+    schemaName?: string,
+    schemaDescription?: string
+  }
+): Promise<Array<z.infer<TSchema>>>;
 
 /**
- * Use a LLM Model to generate object
- *
- * This function a wrapper over AI SDK's generateObject function.
- *
- * It accepts the same arguments of the the original function, plus a "Prompt" object, generated using the `output.ai/prompt`.
+ * Use an LLM model to generate a result from an enum (array of string values).
  *
- * The Prompt object will set `model`, `messages`, `temperature` and `max_tokens`, however all these can be overwritten by their AI SDK native argument values.
+ * This function is a wrapper over the AI SDK's `generateObject` with `output: 'enum'`.
+ * The `Prompt` sets `model`, `messages`, `temperature`, `max_tokens`, and `provider_options`.
  *
- * @param {GenerateObjectArgs} args - Generation arguments
- * @returns {Promise<object>} An object without a pre-defined schema schema
+ * @param {object} args - Generation arguments
+ * @param {Prompt} args.prompt - `@output.ai/prompt` Prompt object
+ * @param {string[]} args.enum - Allowed values for the generation
+ * @returns {Promise<string>} One of the provided enum values
  */
-export function generateObject(
-  args: GenerateObjectArgs
-): Promise<object>;
+export function generateEnum<const TEnum extends readonly [string, ...string[]]>(
+  args: {
+    prompt: Prompt,
+    enum: TEnum
+  }
+): Promise<TEnum[number]>;

package/src/index.js

@@ -1 +1,2 @@
-export { generateText, generateObject } from './ai_sdk.js';
+export { generateText, generateArray, generateObject, generateEnum } from './ai_sdk.js';
+export * as ai from 'ai';

package/src/validations.js

@@ -0,0 +1,52 @@
+import { ValidationError, z } from '@output.ai/core';
+import { promptSchema } from '@output.ai/prompt';
+
+const generateTextArgsSchema = z.object( {
+  prompt: promptSchema
+} );
+
+const generateObjectArgsSchema = z.object( {
+  prompt: promptSchema,
+  schema: z.custom( v => v instanceof z.ZodObject, {
+    message: 'schema must be a ZodObject'
+  } ),
+  schemaName: z.string().optional(),
+  schemaDescription: z.string().optional()
+} );
+
+const generateArrayArgsSchema = z.object( {
+  prompt: promptSchema,
+  schema: z.custom( v => v instanceof z.ZodType, {
+    message: 'schema must be a ZodType'
+  } ),
+  schemaName: z.string().optional(),
+  schemaDescription: z.string().optional()
+} );
+
+const generateEnumArgsSchema = z.object( {
+  prompt: promptSchema,
+  enum: z.array( z.string() )
+} );
+
+export function validateSchema( schema, input, errorPrefix ) {
+  const result = schema.safeParse( input );
+  if ( !result.success ) {
+    throw new ValidationError( `${errorPrefix}: ${z.prettifyError( result.error )}` );
+  }
+};
+
+export function validateGenerateTextArgs( args ) {
+  validateSchema( generateTextArgsSchema, args, 'Invalid generateText() arguments' );
+};
+
+export function validateGenerateObjectArgs( args ) {
+  validateSchema( generateObjectArgsSchema, args, 'Invalid generateObject() arguments' );
+};
+
+export function validateGenerateArrayArgs( args ) {
+  validateSchema( generateArrayArgsSchema, args, 'Invalid generateArray() arguments' );
+};
+
+export function validateGenerateEnumArgs( args ) {
+  validateSchema( generateEnumArgsSchema, args, 'Invalid generateEnum() arguments' );
+};