UNPKG

ai-functions

Version:

Core AI primitives for building intelligent applications

969 lines 35.6 kB
/** * AIPromise - Promise pipelining for AI functions * * This enables: * - Property access tracking for dynamic schema inference * - Promise pipelining without await * - Magical .map() for batch processing * - Dependency graph resolution * * @example * ```ts * // Dynamic schema from destructuring * const { summary, keyPoints, conclusion } = ai`write about ${topic}` * * // Pipeline without await * const isValid = is`${conclusion} is solid given ${keyPoints}` * * // Batch process with map * const ideas = list`startup ideas` * const evaluated = await ideas.map(idea => ({ * idea, * viable: is`${idea} is viable`, * market: ai`market size for ${idea}`, * })) * * // Only await at the end * if (await isValid) { ... } * ``` * * @packageDocumentation */ import { generateObject, streamObject, streamText } from './generate.js'; import { isInRecordingMode, getCurrentItemPlaceholder, captureOperation, createBatchMap, BatchMapPromise, } from './batch-map.js'; import { getModel } from './context.js'; // ============================================================================ // Types // ============================================================================ /** Symbol to identify AIPromise instances */ export const AI_PROMISE_SYMBOL = Symbol.for('ai-promise'); /** Symbol to get the raw AIPromise from a proxy */ export const RAW_PROMISE_SYMBOL = Symbol.for('ai-promise-raw'); /** Recording mode for map() */ export const RECORDING_MODE = Symbol.for('ai-promise-recording'); // ============================================================================ // Global State // ============================================================================ /** Current recording context for map() */ let currentRecording = null; /** Pending promises for batch resolution */ const pendingPromises = new Set(); /** Promise resolution queue */ let resolutionScheduled = false; // ============================================================================ // AIPromise Implementation // ============================================================================ /** * AIPromise - Promise wrapper for AI functions * * Acts as both a Promise AND a stub that: * - Tracks property accesses for dynamic schema inference * - Records dependencies for promise pipelining * - Supports .map() for batch processing */ export class AIPromise { /** Marker to identify AIPromise instances */ [AI_PROMISE_SYMBOL] = true; /** The prompt that will generate this value */ _prompt; /** Options for generation */ _options; /** Properties accessed on this promise (for schema inference) */ _accessedProps = new Set(); /** Property path from parent (for nested access) */ _propertyPath; /** Parent promise (if this is a property access) */ _parent; /** Dependencies (other AIPromises used in our prompt) */ _dependencies = []; /** Cached resolver promise */ _resolver = null; /** Resolved value (cached after first resolution) */ _resolvedValue; /** Whether this promise has been resolved */ _isResolved = false; /** Whether we're in recording mode */ _isRecording = false; constructor(prompt, options = {}) { this._prompt = prompt; this._options = options; this._propertyPath = options.propertyPath || []; this._parent = options.parent || null; // Track this promise for batch resolution pendingPromises.add(this); // Return a proxy that intercepts property access return new Proxy(this, PROXY_HANDLERS); } /** Get the prompt */ get prompt() { return this._prompt; } /** Get the property path */ get path() { return this._propertyPath; } /** Check if resolved */ get isResolved() { return this._isResolved; } /** Get accessed properties */ get accessedProps() { return this._accessedProps; } /** * Add a dependency (another AIPromise used in this one's prompt) */ addDependency(promise, path = []) { this._dependencies.push({ promise, path }); } /** * Resolve this promise */ async resolve() { if (this._isResolved) { return this._resolvedValue; } // If this is a property access on a parent, resolve the parent first if (this._parent) { const parentValue = await this._parent.resolve(); const value = getNestedValue(parentValue, this._propertyPath); this._resolvedValue = value; this._isResolved = true; return this._resolvedValue; } // Resolve dependencies first const resolvedDeps = {}; for (const dep of this._dependencies) { const value = await dep.promise.resolve(); const key = dep.path.length > 0 ? dep.path.join('.') : `dep_${this._dependencies.indexOf(dep)}`; resolvedDeps[key] = value; } // Substitute resolved dependencies into prompt let finalPrompt = this._prompt; for (const [key, value] of Object.entries(resolvedDeps)) { finalPrompt = finalPrompt.replace(new RegExp(`\\$\\{${key}\\}`, 'g'), String(value)); } // Build schema from accessed properties const schema = this._buildSchema(); // Generate the result const result = await generateObject({ model: this._options.model || 'sonnet', schema, prompt: finalPrompt, ...(this._options.system !== undefined && { system: this._options.system }), ...(this._options.temperature !== undefined && { temperature: this._options.temperature }), ...(this._options.maxTokens !== undefined && { maxTokens: this._options.maxTokens }), }); // 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; } /** * Build schema from accessed properties and base schema */ _buildSchema() { const baseSchema = this._options.baseSchema || {}; // If no properties accessed, use base schema or infer from type if (this._accessedProps.size === 0) { if (typeof baseSchema === 'object' && Object.keys(baseSchema).length > 0) { return baseSchema; } // Infer from type 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': return { answer: 'true | false' }; case 'text': return { text: 'The generated text' }; default: return { result: 'The result' }; } } // Build schema from accessed properties const schema = {}; for (const prop of this._accessedProps) { // Check if base schema has this property if (typeof baseSchema === 'object' && !Array.isArray(baseSchema) && prop in baseSchema) { const propSchema = baseSchema[prop]; if (propSchema !== undefined) { schema[prop] = propSchema; continue; } } // Infer type from property name patterns const lowerProp = prop.toLowerCase(); if (lowerProp.endsWith('s') || lowerProp.includes('list') || lowerProp.includes('items') || lowerProp.includes('array')) { schema[prop] = [`List of ${prop}`]; } else if (lowerProp.includes('is') || lowerProp.includes('has') || lowerProp.includes('can') || lowerProp.includes('should')) { schema[prop] = `Whether ${prop} (true/false)`; } else if (lowerProp.includes('count') || lowerProp.includes('number') || lowerProp.includes('total') || lowerProp.includes('amount')) { schema[prop] = `The ${prop} (number)`; } else { schema[prop] = `The ${prop}`; } } return schema; } /** * Map over array results - automatically batches operations! * * When you map over a list, the operations are captured and * automatically batched when resolved. Uses provider batch APIs * for cost savings (50% discount) when beneficial. * * @example * ```ts * // Simple map - each title becomes a blog post * const titles = await list`10 blog post titles` * const posts = titles.map(title => write`blog post: # ${title}`) * console.log(await posts) // 10 blog posts via batch API * * // Complex map - multiple operations per item * const ideas = await list`startup ideas` * const evaluated = await ideas.map(idea => ({ * idea, * viable: is`${idea} is viable`, * market: ai`market size for ${idea}`, * })) * ``` */ map(callback) { // Create a wrapper that resolves this promise first, then maps const mapPromise = new BatchMapPromise([], [], {}); // Override the resolve to first get the list items // 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; } /** * Map with explicit batch options * * @example * ```ts * // Force immediate execution (no batch API) * const posts = titles.mapImmediate(title => write`blog post: ${title}`) * * // Force batch API (even for small lists) * const posts = titles.mapDeferred(title => write`blog post: ${title}`) * ``` */ mapImmediate(callback) { const mapPromise = new BatchMapPromise([], [], { immediate: true }); 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 }); 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; } /** * ForEach with automatic batching * * @example * ```ts * await list`startup ideas`.forEach(async idea => { * console.log(await is`${idea} is viable`) * }) * ``` */ async forEach(callback) { const items = await this.resolve(); if (Array.isArray(items)) { for (let i = 0; i < items.length; i++) { await callback(items[i], i); } } else { // When T is not an array, the conditional type T extends (infer I)[] ? I : T resolves to T await callback(items, 0); } } /** * Async iterator support with smart batching */ async *[Symbol.asyncIterator]() { 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() */ then(onfulfilled, onrejected) { if (!this._resolver) { // Schedule batch resolution on next microtask this._resolver = new Promise((resolve, reject) => { queueMicrotask(async () => { try { const value = await this.resolve(); resolve(value); } catch (error) { reject(error); } }); }); } return this._resolver.then(onfulfilled, onrejected); } /** * Promise interface - catch() */ catch(onrejected) { return this.then(null, onrejected); } /** * Promise interface - finally() */ finally(onfinally) { return this.then((value) => { onfinally?.(); return value; }, (reason) => { onfinally?.(); throw reason; }); } } // ============================================================================ // Proxy Handlers // ============================================================================ const PROXY_HANDLERS = { get(target, prop, _receiver) { // Handle symbols if (typeof prop === 'symbol') { if (prop === AI_PROMISE_SYMBOL) return true; if (prop === RAW_PROMISE_SYMBOL) return target; if (prop === Symbol.asyncIterator) return target[Symbol.asyncIterator].bind(target); return target[prop]; } // Handle promise methods if (prop === 'then' || prop === 'catch' || prop === 'finally') { const method = target[prop]; return method?.bind(target); } // Handle AIPromise methods 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') { return target[prop]; } // Track property access for schema inference target.accessedProps.add(prop); // If we're in recording mode, record this access if (currentRecording) { // Just track the access, don't create new promise } // Return a new AIPromise for the property path return new AIPromise(target.prompt, { ...target['_options'], parent: target, propertyPath: [...target.path, prop], }); }, // Prevent mutation set() { throw new Error('AIPromise properties are read-only'); }, deleteProperty() { throw new Error('AIPromise properties cannot be deleted'); }, // Handle function calls (for chained methods) apply(target, _thisArg, args) { // If the target is callable (e.g., from a template function), call it const call = target['_call']; if (typeof call === 'function') { return call(...args); } throw new Error('AIPromise is not callable'); }, }; // ============================================================================ // Helper Functions // ============================================================================ /** * Get a nested value from an object by path */ function getNestedValue(obj, path) { let current = obj; for (const key of path) { if (current === null || current === undefined) return undefined; if (key === '__item__') continue; // Skip internal markers current = current[key]; } return current; } /** * Analyze the result of a map callback to build batch schema */ function analyzeRecordingResult(result, recording) { if (result === null || result === undefined) { return { result: 'The result' }; } if (typeof result !== 'object') { return { result: 'The result' }; } // Build schema from the result structure const schema = {}; for (const [key, value] of Object.entries(result)) { if (isAIPromise(value)) { // This is a reference to an AI operation 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}`])); } else { // Access private _options through type-safe assertion const options = aiPromise._options; const type = options?.type; if (type === 'boolean') { schema[key] = 'true | false'; } else if (type === 'list') { schema[key] = ['List items']; } else { schema[key] = `The ${key}`; } } } else if (typeof value === 'object' && value !== null) { // Recursively analyze nested objects schema[key] = analyzeRecordingResult(value, recording); } else { // Literal value - include as-is schema[key] = `Value: ${JSON.stringify(value)}`; } } return schema; } /** * Check if a value is an AIPromise */ export function isAIPromise(value) { return (value !== null && typeof value === 'object' && AI_PROMISE_SYMBOL in value && value[AI_PROMISE_SYMBOL] === true); } /** * Get the raw AIPromise from a proxied value */ export function getRawPromise(value) { const raw = value[RAW_PROMISE_SYMBOL]; return raw ?? value; } // ============================================================================ // Factory Functions // ============================================================================ /** * Create an AIPromise for text generation */ export function createTextPromise(prompt, options) { return new AIPromise(prompt, { ...options, type: 'text' }); } /** * Create an AIPromise for object generation with dynamic schema */ export function createObjectPromise(prompt, options) { return new AIPromise(prompt, { ...options, type: 'object' }); } /** * Create an AIPromise for list generation */ export function createListPromise(prompt, options) { return new AIPromise(prompt, { ...options, type: 'list' }); } /** * Create an AIPromise for multiple lists generation */ export function createListsPromise(prompt, options) { return new AIPromise(prompt, { ...options, type: 'lists' }); } /** * Create an AIPromise for boolean/is check */ export function createBooleanPromise(prompt, options) { return new AIPromise(prompt, { ...options, type: 'boolean' }); } /** * Create an AIPromise for extraction */ export function createExtractPromise(prompt, options) { return new AIPromise(prompt, { ...options, type: 'extract' }); } // ============================================================================ // Template Tag Helpers // ============================================================================ /** * Parse template literals and track AIPromise dependencies */ export function parseTemplateWithDependencies(strings, ...values) { const dependencies = []; let prompt = ''; for (let i = 0; i < strings.length; i++) { prompt += strings[i]; if (i < values.length) { const value = values[i]; if (isAIPromise(value)) { // Track as dependency const rawPromise = getRawPromise(value); const depKey = `dep_${dependencies.length}`; dependencies.push({ promise: rawPromise, path: rawPromise.path }); prompt += `\${${depKey}}`; } else { // Inline the value prompt += String(value); } } } return { prompt, dependencies }; } /** * Create a template function that returns AIPromise */ export function createAITemplateFunction(type, baseOptions) { function templateFn(promptOrStrings, ...args) { let prompt; let dependencies = []; let options = { ...baseOptions }; if (Array.isArray(promptOrStrings) && 'raw' in promptOrStrings) { // Tagged template literal const parsed = parseTemplateWithDependencies(promptOrStrings, ...args); prompt = parsed.prompt; dependencies = parsed.dependencies; } else { // Regular function call prompt = promptOrStrings; if (args.length > 0 && typeof args[0] === 'object') { options = { ...options, ...args[0] }; } } // 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'; captureOperation(prompt, batchType, options.baseSchema, options.system); } 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