@output.ai/llm 0.0.12 → 0.0.14

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@output.ai/llm",
-  "version": "0.0.12",
+  "version": "0.0.14",
   "description": "Framework abstraction to interact with LLM models",
   "type": "module",
   "main": "src/index.js",
@@ -9,10 +9,11 @@
     "./src"
   ],
   "dependencies": {
-    "@ai-sdk/anthropic": "2.0.4",
-    "@ai-sdk/azure": "2.0.49",
-    "@ai-sdk/openai": "2.0.18",
-    "@output.ai/trace": "0.0.1",
+    "@ai-sdk/anthropic": "2.0.28",
+    "@ai-sdk/azure": "2.0.53",
+    "@ai-sdk/openai": "2.0.52",
+    "@output.ai/core": ">=0.0.1",
+    "@output.ai/prompt": ">=0.0.1",
     "ai": "5.0.48"
   },
   "license": "UNLICENSED"

package/src/ai_model.js

@@ -0,0 +1,29 @@
+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_sdk.js

@@ -1,73 +1,45 @@
-import {
-  generateText as aiGenerateText,
-  generateObject as aiGenerateObject
-} from 'ai';
-import { anthropic } from '@ai-sdk/anthropic';
-import { openai } from '@ai-sdk/openai';
-import { azure } from '@ai-sdk/azure';
-import { trace } from '@output.ai/trace';
-
-const providers = {
-  anthropic,
-  azure,
-  openai
-};
-
-export async function generateText( options ) {
-  const { prompt, ...aiSdkOptions } = options;
-
-  const provider = providers[prompt.config.provider];
-  const model = provider( prompt.config.model );
-
-  const aiOptions = {
-    model,
-    messages: prompt.messages,
-    temperature: prompt.config.temperature,
-    ...aiSdkOptions // Spread remaining AI SDK options (thinking, etc.)
-  };
-
-  if ( prompt.config.max_tokens ) {
-    aiOptions.maxOutputTokens = prompt.config.max_tokens;
+import { Tracing } from '@output.ai/core/tracing';
+import { loadModel } from './ai_model.js';
+import * as AI from 'ai';
+
+const generationWrapper = async ( traceId, fn ) => {
+  try {
+    const result = await fn();
+    Tracing.addEventEnd( { id: traceId, details: result } );
+    return result;
+  } catch ( error ) {
+    Tracing.addEventError( { id: traceId, details: error } );
+    throw error;
   }
+};
 
-  const result = await aiGenerateText( aiOptions );
-
-  trace( { lib: 'llm', event: 'generateText', input: options, output: result } );
-
-  return result.text;
+export async function generateText( { prompt, ...nativeAiSdkArgs } ) {
+  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;
+  } );
 }
 
-export async function generateObject( options ) {
-  const {
-    prompt,
-    schema,
-    schemaName,
-    schemaDescription,
-    output = 'object',
-    ...aiSdkOptions
-  } = options;
-
-  const provider = providers[prompt.config.provider];
-  const model = provider( prompt.config.model );
-
-  const aiOptions = {
-    model,
-    schema,
-    schemaName,
-    schemaDescription,
-    output,
-    messages: prompt.messages,
-    temperature: prompt.config.temperature,
-    ...aiSdkOptions // Spread remaining AI SDK options
-  };
-
-  if ( prompt.config.max_tokens ) {
-    aiOptions.maxOutputTokens = prompt.config.max_tokens;
-  }
-
-  const result = await aiGenerateObject( aiOptions );
-
-  trace( { lib: 'llm', event: 'generateObject', input: options, output: result } );
-
-  return result.object;
+export async function generateObject( { prompt, ...nativeAiSdkArgs } ) {
+  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;
+  } );
 }

package/src/index.d.ts

@@ -1,44 +1,60 @@
-import type {
-  GenerateTextOptions,
-  GenerateObjectOptions
-} from 'ai';
+import type * as AiTypes from 'ai';
+import type { Prompt } from '@output.ai/prompt';
 
-export interface Prompt {
-  config: {
-    provider: 'anthropic' | 'openai' | 'azure' ;
-    model: string;
-    temperature?: number;
-    max_tokens?: number;
-  };
-  messages: Array<{
-    role: string;
-    content: string;
-  }>;
-}
+export type { Prompt };
 
-// Omit prompt/messages from AI SDK options and add our Prompt type
-export type GenerateTextOptionsWithPrompt = Omit<
-  GenerateTextOptions,
-  'prompt' | 'messages' | 'model'
-> & {
-  prompt: Prompt,
-};
+type NativeGenerateTextArgs = Parameters<typeof AiTypes.generateText>[0];
+type NativeGenerateObjectArgs = Parameters<typeof AiTypes.generateObject>[0];
 
-export type GenerateObjectOptionsWithPrompt<SCHEMA> = Omit<
-  GenerateObjectOptions<SCHEMA>,
-  'prompt' | 'messages' | 'model' | 'schema' | 'schemaName' | 'schemaDescription' | 'output'
-> & {
-  prompt: Prompt,
-  schema: Record<string, unknown> | object,
-  schemaName?: string,
-  schemaDescription?: string,
-  output?: 'object' | 'array',
-};
+/**
+ * Simplify types into a plain object
+ */
+type Simplify<T> = { [K in keyof T]: T[K] } & {};
 
+/**
+ * Text generation arguments
+ * Include all native AI SDK generateText options and a Prompt object from `@output.ai/prompt`
+ */
+export type GenerateTextArgs = Simplify<
+  Partial<Omit<NativeGenerateTextArgs, 'prompt'>> & { prompt: Prompt }
+>;
+
+/**
+ * Object generation arguments
+ * Include all native AI SDK generateObject options and a Prompt object from `@output.ai/prompt`
+ */
+export type GenerateObjectArgs = Simplify<
+  Partial<Omit<NativeGenerateObjectArgs, 'prompt'>> & { 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`.
+ *
+ * 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 {GenerateTextArgs} args - Generation arguments
+ * @returns {Promise<string>}
+ */
 export function generateText(
-  options: GenerateTextOptionsWithPrompt
+  args: GenerateTextArgs
 ): Promise<string>;
 
-export function generateObject<T = unknown>(
-  options: GenerateObjectOptionsWithPrompt<T>
-): Promise<T>;
+/**
+ * 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`.
+ *
+ * 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>}
+ */
+export function generateObject(
+  args: GenerateObjectArgs
+): Promise<object>;