@output.ai/llm 0.0.14 → 0.0.15

package/package.json

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

package/src/index.d.ts

@@ -1,4 +1,5 @@
 import type * as AiTypes from 'ai';
+import type { z as CoreZ } from '@output.ai/core';
 import type { Prompt } from '@output.ai/prompt';
 
 export type { Prompt };
@@ -7,24 +8,41 @@ type NativeGenerateTextArgs = Parameters<typeof AiTypes.generateText>[0];
 type NativeGenerateObjectArgs = Parameters<typeof AiTypes.generateObject>[0];
 
 /**
- * Simplify types into a plain object
+ * Simplify types into a plain object while preserving unions
+ * (distributes over union members instead of collapsing their keys)
  */
-type Simplify<T> = { [K in keyof T]: T[K] } & {};
+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<
-  Partial<Omit<NativeGenerateTextArgs, 'prompt'>> & { prompt: Prompt }
+  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<
-  Partial<Omit<NativeGenerateObjectArgs, 'prompt'>> & { prompt: Prompt }
+  Replace<Partial<WithCoreZodSchema<NativeGenerateObjectArgs>>, 'prompt', Prompt>
 >;
 
 /**
@@ -53,7 +71,39 @@ export function generateText(
  * 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>}
+ * @returns {Promise<object>} An object matching the provided schema
+ */
+export function generateObject<A extends GenerateObjectArgs & { schema: CoreZ.ZodTypeAny }>(
+  args: A
+): Promise<CoreZ.infer<A['schema']>>;
+
+/**
+ * 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>} An object matching the provided enum
+ */
+export function generateObject<A extends GenerateObjectArgs & { enum: readonly unknown[]; output: 'enum' }>(
+  args: A
+): Promise<A['enum'][number]>;
+
+/**
+ * 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>} An object without a pre-defined schema schema
  */
 export function generateObject(
   args: GenerateObjectArgs