ai-functions 2.1.1 → 2.3.0

package/dist/ai-promise.js

@@ -1,7 +1,7 @@
 /**
- * AIPromise - RPC-style promise pipelining for AI functions
+ * AIPromise - Promise pipelining for AI functions
  *
- * Inspired by capnweb's RpcPromise, this enables:
+ * This enables:
  * - Property access tracking for dynamic schema inference
  * - Promise pipelining without await
  * - Magical .map() for batch processing
@@ -29,7 +29,7 @@
  *
  * @packageDocumentation
  */
-import { generateObject } from './generate.js';
+import { generateObject, streamObject, streamText } from './generate.js';
 import { isInRecordingMode, getCurrentItemPlaceholder, captureOperation, createBatchMap, BatchMapPromise, } from './batch-map.js';
 import { getModel } from './context.js';
 // ============================================================================
@@ -54,7 +54,7 @@ let resolutionScheduled = false;
 // AIPromise Implementation
 // ============================================================================
 /**
- * AIPromise - Like capnweb's RpcPromise but for AI functions
+ * AIPromise - Promise wrapper for AI functions
  *
  * Acts as both a Promise AND a stub that:
  * - Tracks property accesses for dynamic schema inference
@@ -150,11 +150,39 @@ export class AIPromise {
             model: this._options.model || 'sonnet',
             schema,
             prompt: finalPrompt,
-            system: this._options.system,
-            temperature: this._options.temperature,
-            maxTokens: this._options.maxTokens,
+            ...(this._options.system !== undefined && { system: this._options.system }),
+            ...(this._options.temperature !== undefined && { temperature: this._options.temperature }),
+            ...(this._options.maxTokens !== undefined && { maxTokens: this._options.maxTokens }),
         });
-        this._resolvedValue = result.object;
+        // Extract the value based on type
+        // Type assertions here are safe because:
+        // 1. Runtime type checking validates the response structure
+        // 2. The type parameter T corresponds to the expected output type for each mode
+        let value = result.object;
+        if (this._options.type === 'text' &&
+            typeof value === 'object' &&
+            value !== null &&
+            'text' in value) {
+            value = value.text;
+        }
+        else if (this._options.type === 'boolean' &&
+            typeof value === 'object' &&
+            value !== null &&
+            'answer' in value) {
+            const answer = value.answer;
+            // When type === 'boolean', T is constrained to boolean at the call site.
+            // TypeScript can't express this dependent relationship, so we use a simple cast.
+            // Runtime validation: answer is verified to be 'true', 'false', or boolean.
+            const booleanValue = answer === 'true' || answer === true;
+            value = booleanValue;
+        }
+        else if ((this._options.type === 'list' || this._options.type === 'extract') &&
+            typeof value === 'object' &&
+            value !== null &&
+            'items' in value) {
+            value = value.items;
+        }
+        this._resolvedValue = value;
         this._isResolved = true;
         pendingPromises.delete(this);
         return this._resolvedValue;
@@ -173,6 +201,12 @@ export class AIPromise {
             switch (this._options.type) {
                 case 'list':
                     return { items: ['List items'] };
+                case 'extract':
+                    return {
+                        items: [
+                            'Array of extracted items as strings - extract ALL matching items from the text',
+                        ],
+                    };
                 case 'lists':
                     return { categories: ['Category names'], data: 'JSON object with categorized lists' };
                 case 'boolean':
@@ -247,17 +281,23 @@ export class AIPromise {
         // Create a wrapper that resolves this promise first, then maps
         const mapPromise = new BatchMapPromise([], [], {});
         // Override the resolve to first get the list items
-        const originalResolve = mapPromise.resolve.bind(mapPromise);
-        mapPromise.resolve = async () => {
-            // First, resolve the list
-            const items = await this.resolve();
-            if (!Array.isArray(items)) {
-                throw new Error('Cannot map over non-array result');
-            }
-            // Now create the actual batch map with the resolved items
-            const actualBatchMap = createBatchMap(items, callback);
-            return actualBatchMap.resolve();
-        };
+        // Type assertion: BatchMapPromise.resolve is a public method that we're replacing
+        // with a compatible async function returning Promise<U[]>
+        const self = this;
+        Object.defineProperty(mapPromise, 'resolve', {
+            value: async function () {
+                // First, resolve the list
+                const items = await self.resolve();
+                if (!Array.isArray(items)) {
+                    throw new Error('Cannot map over non-array result');
+                }
+                // Now create the actual batch map with the resolved items
+                const actualBatchMap = createBatchMap(items, callback);
+                return actualBatchMap.resolve();
+            },
+            writable: true,
+            configurable: true,
+        });
         return mapPromise;
     }
     /**
@@ -274,26 +314,36 @@ export class AIPromise {
      */
     mapImmediate(callback) {
         const mapPromise = new BatchMapPromise([], [], { immediate: true });
-        mapPromise.resolve = async () => {
-            const items = await this.resolve();
-            if (!Array.isArray(items)) {
-                throw new Error('Cannot map over non-array result');
-            }
-            const actualBatchMap = createBatchMap(items, callback, { immediate: true });
-            return actualBatchMap.resolve();
-        };
+        const self = this;
+        Object.defineProperty(mapPromise, 'resolve', {
+            value: async function () {
+                const items = await self.resolve();
+                if (!Array.isArray(items)) {
+                    throw new Error('Cannot map over non-array result');
+                }
+                const actualBatchMap = createBatchMap(items, callback, { immediate: true });
+                return actualBatchMap.resolve();
+            },
+            writable: true,
+            configurable: true,
+        });
         return mapPromise;
     }
     mapDeferred(callback) {
         const mapPromise = new BatchMapPromise([], [], { deferred: true });
-        mapPromise.resolve = async () => {
-            const items = await this.resolve();
-            if (!Array.isArray(items)) {
-                throw new Error('Cannot map over non-array result');
-            }
-            const actualBatchMap = createBatchMap(items, callback, { deferred: true });
-            return actualBatchMap.resolve();
-        };
+        const self = this;
+        Object.defineProperty(mapPromise, 'resolve', {
+            value: async function () {
+                const items = await self.resolve();
+                if (!Array.isArray(items)) {
+                    throw new Error('Cannot map over non-array result');
+                }
+                const actualBatchMap = createBatchMap(items, callback, { deferred: true });
+                return actualBatchMap.resolve();
+            },
+            writable: true,
+            configurable: true,
+        });
         return mapPromise;
     }
     /**
@@ -314,6 +364,7 @@ export class AIPromise {
             }
         }
         else {
+            // When T is not an array, the conditional type T extends (infer I)[] ? I : T resolves to T
             await callback(items, 0);
         }
     }
@@ -324,13 +375,43 @@ export class AIPromise {
         const items = await this.resolve();
         if (Array.isArray(items)) {
             for (const item of items) {
+                // Each array item is the inferred element type I when T extends I[]
                 yield item;
             }
         }
         else {
+            // When T is not an array, the item type is T itself
             yield items;
         }
     }
+    /**
+     * Stream the AI generation - returns chunks as they arrive
+     *
+     * For text generation, yields string chunks.
+     * For object generation, yields partial objects as they build up.
+     * For list generation, yields items as they're generated.
+     *
+     * @example
+     * ```ts
+     * // Text streaming
+     * const stream = write`Write a story`.stream()
+     * for await (const chunk of stream.textStream) {
+     *   process.stdout.write(chunk)
+     * }
+     *
+     * // Object streaming with partial updates
+     * const stream = ai`Generate a recipe`.stream()
+     * for await (const partial of stream.partialObjectStream) {
+     *   console.log('Building:', partial)
+     * }
+     *
+     * // Get final result after streaming
+     * const finalResult = await stream.result
+     * ```
+     */
+    stream(options) {
+        return createStreamingAIPromise(this, options);
+    }
     /**
      * Promise interface - then()
      */
@@ -374,7 +455,7 @@ export class AIPromise {
 // Proxy Handlers
 // ============================================================================
 const PROXY_HANDLERS = {
-    get(target, prop, receiver) {
+    get(target, prop, _receiver) {
         // Handle symbols
         if (typeof prop === 'symbol') {
             if (prop === AI_PROMISE_SYMBOL)
@@ -387,14 +468,26 @@ const PROXY_HANDLERS = {
         }
         // Handle promise methods
         if (prop === 'then' || prop === 'catch' || prop === 'finally') {
-            return target[prop].bind(target);
+            const method = target[prop];
+            return method?.bind(target);
         }
         // Handle AIPromise methods
-        if (prop === 'map' || prop === 'forEach' || prop === 'resolve') {
-            return target[prop].bind(target);
+        if (prop === 'map' ||
+            prop === 'forEach' ||
+            prop === 'resolve' ||
+            prop === 'stream' ||
+            prop === 'addDependency' ||
+            prop === 'mapImmediate' ||
+            prop === 'mapDeferred') {
+            const method = target[prop];
+            return method?.bind(target);
         }
         // Handle internal properties
-        if (prop.startsWith('_') || prop === 'prompt' || prop === 'path' || prop === 'isResolved' || prop === 'accessedProps') {
+        if (prop.startsWith('_') ||
+            prop === 'prompt' ||
+            prop === 'path' ||
+            prop === 'isResolved' ||
+            prop === 'accessedProps') {
             return target[prop];
         }
         // Track property access for schema inference
@@ -418,10 +511,11 @@ const PROXY_HANDLERS = {
         throw new Error('AIPromise properties cannot be deleted');
     },
     // Handle function calls (for chained methods)
-    apply(target, thisArg, args) {
+    apply(target, _thisArg, args) {
         // If the target is callable (e.g., from a template function), call it
-        if (typeof target._call === 'function') {
-            return target._call(...args);
+        const call = target['_call'];
+        if (typeof call === 'function') {
+            return call(...args);
         }
         throw new Error('AIPromise is not callable');
     },
@@ -461,10 +555,12 @@ function analyzeRecordingResult(result, recording) {
             const aiPromise = getRawPromise(value);
             // Infer schema from the promise's accessed properties or type
             if (aiPromise.accessedProps.size > 0) {
-                schema[key] = Object.fromEntries(Array.from(aiPromise.accessedProps).map(p => [p, `The ${p}`]));
+                schema[key] = Object.fromEntries(Array.from(aiPromise.accessedProps).map((p) => [p, `The ${p}`]));
             }
             else {
-                const type = aiPromise._options?.type;
+                // Access private _options through type-safe assertion
+                const options = aiPromise._options;
+                const type = options?.type;
                 if (type === 'boolean') {
                     schema[key] = 'true | false';
                 }
@@ -500,10 +596,8 @@ export function isAIPromise(value) {
  * Get the raw AIPromise from a proxied value
  */
 export function getRawPromise(value) {
-    if (RAW_PROMISE_SYMBOL in value) {
-        return value[RAW_PROMISE_SYMBOL];
-    }
-    return value;
+    const raw = value[RAW_PROMISE_SYMBOL];
+    return raw ?? value;
 }
 // ============================================================================
 // Factory Functions
@@ -595,16 +689,281 @@ export function createAITemplateFunction(type, baseOptions) {
         }
         // If we're in recording mode (inside a .map() callback), capture this operation
         if (isInRecordingMode()) {
-            const batchType = type === 'text' ? 'text' : type === 'boolean' ? 'boolean' : type === 'list' ? 'list' : 'object';
+            const batchType = type === 'text'
+                ? 'text'
+                : type === 'boolean'
+                    ? 'boolean'
+                    : type === 'list'
+                        ? 'list'
+                        : 'object';
             captureOperation(prompt, batchType, options.baseSchema, options.system);
         }
-        const promise = new AIPromise(prompt, { ...options, type });
+        const promise = new AIPromise(prompt, {
+            ...options,
+            ...(type !== undefined && { type }),
+        });
         // Add dependencies
         for (const dep of dependencies) {
             promise.addDependency(dep.promise, dep.path);
         }
         return promise;
     }
+    // Return type matches the declared intersection type
     return templateFn;
 }
+// ============================================================================
+// Streaming Implementation
+// ============================================================================
+/**
+ * Create a streaming wrapper for an AIPromise
+ *
+ * This function creates a StreamingAIPromise that:
+ * - Resolves dependencies before streaming
+ * - Streams text or partial objects based on the promise type
+ * - Collects the final result as stream is consumed
+ * - Supports cancellation via AbortSignal
+ */
+function createStreamingAIPromise(promise, options) {
+    const rawPromise = getRawPromise(promise);
+    const promiseOptions = rawPromise._options;
+    const dependencies = rawPromise._dependencies;
+    // Result promise state
+    let resultResolve;
+    let resultReject;
+    const resultPromise = new Promise((resolve, reject) => {
+        resultResolve = resolve;
+        resultReject = reject;
+    });
+    // Shared state to prevent multiple API calls
+    let streamStarted = false;
+    let cachedTextChunks = null;
+    let cachedPartialObjects = null;
+    let streamError = null;
+    let finalValue;
+    // Resolve dependencies and prepare the final prompt
+    const preparePrompt = async () => {
+        const resolvedDeps = {};
+        for (const dep of dependencies) {
+            const value = await dep.promise.resolve();
+            const key = dep.path.length > 0 ? dep.path.join('.') : `dep_${dependencies.indexOf(dep)}`;
+            resolvedDeps[key] = value;
+        }
+        let finalPrompt = rawPromise.prompt;
+        for (const [key, value] of Object.entries(resolvedDeps)) {
+            finalPrompt = finalPrompt.replace(new RegExp(`\\$\\{${key}\\}`, 'g'), String(value));
+        }
+        return finalPrompt;
+    };
+    // Build schema from accessed properties
+    const buildSchema = () => {
+        return rawPromise._buildSchema();
+    };
+    // Extract value based on type (same logic as resolve())
+    // Type assertions here are safe because:
+    // 1. Runtime type checking validates the response structure
+    // 2. The type parameter T corresponds to the expected output type for each mode
+    const extractFinalValue = (obj) => {
+        let value = obj;
+        if (promiseOptions.type === 'text' &&
+            typeof value === 'object' &&
+            value !== null &&
+            'text' in value) {
+            value = value.text;
+        }
+        else if (promiseOptions.type === 'boolean' &&
+            typeof value === 'object' &&
+            value !== null &&
+            'answer' in value) {
+            const answer = value.answer;
+            // When type === 'boolean', T is constrained to boolean at the call site.
+            // TypeScript can't express this dependent relationship, so we use a simple cast.
+            // Runtime validation: answer is verified to be 'true', 'false', or boolean.
+            const booleanValue = answer === 'true' || answer === true;
+            value = booleanValue;
+        }
+        else if ((promiseOptions.type === 'list' || promiseOptions.type === 'extract') &&
+            typeof value === 'object' &&
+            value !== null &&
+            'items' in value) {
+            value = value.items;
+        }
+        return value;
+    };
+    // Create text stream that collects chunks for result
+    async function* createTextStream() {
+        if (cachedTextChunks !== null) {
+            // Return cached chunks if we already streamed
+            for (const chunk of cachedTextChunks) {
+                yield chunk;
+            }
+            return;
+        }
+        if (streamStarted && streamError) {
+            throw streamError;
+        }
+        streamStarted = true;
+        cachedTextChunks = [];
+        try {
+            const finalPrompt = await preparePrompt();
+            const result = await streamText({
+                model: promiseOptions.model || 'sonnet',
+                prompt: finalPrompt,
+                ...(promiseOptions.system !== undefined && { system: promiseOptions.system }),
+                ...(promiseOptions.temperature !== undefined && {
+                    temperature: promiseOptions.temperature,
+                }),
+                ...(promiseOptions.maxTokens !== undefined && { maxTokens: promiseOptions.maxTokens }),
+                ...(options?.abortSignal !== undefined && { abortSignal: options.abortSignal }),
+            });
+            let fullText = '';
+            for await (const chunk of result.textStream) {
+                cachedTextChunks.push(chunk);
+                fullText += chunk;
+                yield chunk;
+            }
+            finalValue = fullText;
+            resultResolve(finalValue);
+        }
+        catch (error) {
+            streamError = error;
+            resultReject(error);
+            throw error;
+        }
+    }
+    // Create partial object stream that collects objects for result
+    async function* createPartialObjectStream() {
+        if (cachedPartialObjects !== null) {
+            // Return cached partials if we already streamed
+            for (const partial of cachedPartialObjects) {
+                yield partial;
+            }
+            return;
+        }
+        if (streamStarted && streamError) {
+            throw streamError;
+        }
+        streamStarted = true;
+        cachedPartialObjects = [];
+        try {
+            const finalPrompt = await preparePrompt();
+            const schema = buildSchema();
+            const result = await streamObject({
+                model: promiseOptions.model || 'sonnet',
+                schema,
+                prompt: finalPrompt,
+                ...(promiseOptions.system !== undefined && { system: promiseOptions.system }),
+                ...(promiseOptions.temperature !== undefined && {
+                    temperature: promiseOptions.temperature,
+                }),
+                ...(promiseOptions.maxTokens !== undefined && { maxTokens: promiseOptions.maxTokens }),
+                ...(options?.abortSignal !== undefined && { abortSignal: options.abortSignal }),
+            });
+            let lastPartial = {};
+            for await (const partial of result.partialObjectStream) {
+                cachedPartialObjects.push(partial);
+                lastPartial = partial;
+                yield partial;
+            }
+            finalValue = extractFinalValue(lastPartial);
+            resultResolve(finalValue);
+        }
+        catch (error) {
+            streamError = error;
+            resultReject(error);
+            throw error;
+        }
+    }
+    // Create main stream based on type
+    async function* createMainStream() {
+        if (promiseOptions.type === 'text') {
+            for await (const chunk of createTextStream()) {
+                // When type is 'text', T is string, so the conditional type resolves to string
+                yield chunk;
+            }
+        }
+        else if (promiseOptions.type === 'list') {
+            // For lists, yield new items as they appear
+            let lastLength = 0;
+            for await (const partial of createPartialObjectStream()) {
+                const items = partial.items || [];
+                for (let i = lastLength; i < items.length; i++) {
+                    // List items are strings, cast to the conditional return type
+                    yield items[i];
+                }
+                lastLength = items.length;
+            }
+        }
+        else {
+            for await (const partial of createPartialObjectStream()) {
+                // For object types, T is not string, so conditional type resolves to Partial<T>
+                yield partial;
+            }
+        }
+    }
+    // Start the stream collection in background if result is awaited
+    const ensureStreamStarted = () => {
+        if (!streamStarted) {
+            // Start consuming the appropriate stream to populate result
+            if (promiseOptions.type === 'text') {
+                ;
+                (async () => {
+                    try {
+                        for await (const _ of createTextStream()) {
+                            // consume
+                        }
+                    }
+                    catch {
+                        // Error already handled in stream
+                    }
+                })();
+            }
+            else {
+                ;
+                (async () => {
+                    try {
+                        for await (const _ of createPartialObjectStream()) {
+                            // consume
+                        }
+                    }
+                    catch {
+                        // Error already handled in stream
+                    }
+                })();
+            }
+        }
+    };
+    // Create a lazy result promise that starts streaming when accessed
+    const lazyResult = Object.assign({
+        then(onfulfilled, onrejected) {
+            ensureStreamStarted();
+            return resultPromise.then(onfulfilled, onrejected);
+        },
+        catch(onrejected) {
+            ensureStreamStarted();
+            return resultPromise.catch(onrejected);
+        },
+        finally(onfinally) {
+            ensureStreamStarted();
+            return resultPromise.finally(onfinally);
+        },
+        [Symbol.toStringTag]: 'Promise',
+    });
+    // Create the streaming object
+    const streamingPromise = {
+        textStream: {
+            [Symbol.asyncIterator]: createTextStream,
+        },
+        partialObjectStream: {
+            [Symbol.asyncIterator]: createPartialObjectStream,
+        },
+        result: lazyResult,
+        [Symbol.asyncIterator]: createMainStream,
+        then(onfulfilled, onrejected) {
+            // If result is awaited before stream consumption, start the stream
+            ensureStreamStarted();
+            return resultPromise.then(onfulfilled, onrejected);
+        },
+    };
+    return streamingPromise;
+}
 //# sourceMappingURL=ai-promise.js.map