ai-functions/dist/batch-queue.js

Core AI primitives for building intelligent applications

/**
 * BatchQueue - Deferred execution queue for batch processing
 *
 * Collects AI operations and submits them to provider batch APIs
 * for cost-effective processing (typically 50% discount, 24hr turnaround).
 *
 * @example
 * ```ts
 * // Create a batch queue
 * const batch = createBatch({ provider: 'openai' })
 *
 * // Add items to the batch (these are deferred, not executed)
 * const titles = await list`10 blog post titles about startups`
 * const posts = titles.map(title => batch.add(write`blog post about ${title}`))
 *
 * // Submit the batch - returns job info
 * const job = await batch.submit()
 * console.log(job.id) // batch_abc123
 *
 * // Poll for results or use webhook
 * const results = await batch.wait()
 * ```
 *
 * @packageDocumentation
 */
import { getLogger } from './logger.js';
import { policyFor } from 'language-models';
/**
 * BatchQueue collects AI operations for deferred batch execution
 */
export class BatchQueue {
    items = [];
    options;
    idCounter = 0;
    submitted = false;
    job = null;
    // Error handling properties
    _autoSubmitPromise = null;
    _submissionError = null;
    _eventHandlers = new Map();
    constructor(options = {}) {
        this.options = {
            provider: 'openai',
            maxItems: 50000, // OpenAI's limit
            autoSubmit: false,
            ...options,
        };
    }
    /**
     * Subscribe to batch events
     * @param event - Event name ('error', 'partial-failure', 'complete')
     * @param handler - Event handler function
     */
    on(event, handler) {
        if (!this._eventHandlers.has(event)) {
            this._eventHandlers.set(event, new Set());
        }
        this._eventHandlers.get(event).add(handler);
    }
    /**
     * Unsubscribe from batch events
     */
    off(event, handler) {
        this._eventHandlers.get(event)?.delete(handler);
    }
    /**
     * Emit an event to all subscribed handlers
     */
    emit(event, data) {
        this._eventHandlers.get(event)?.forEach((handler) => handler(data));
    }
    /**
     * Get the auto-submit promise (if auto-submit was triggered)
     */
    get autoSubmitPromise() {
        return this._autoSubmitPromise ?? undefined;
    }
    /**
     * Get the last submission error (if any)
     */
    get submissionError() {
        return this._submissionError ?? undefined;
    }
    /**
     * Check if there was a submission error
     */
    get hasSubmissionError() {
        return this._submissionError !== null;
    }
    /**
     * Await auto-submit completion or failure
     * Returns a promise that resolves when auto-submit completes or rejects on error
     */
    awaitAutoSubmit() {
        if (!this._autoSubmitPromise) {
            return Promise.resolve();
        }
        return this._autoSubmitPromise;
    }
    /**
     * Get items that failed during batch processing
     */
    getFailedItems() {
        return this.items.filter((item) => item.status === 'failed');
    }
    /**
     * Retry a failed submission
     * Only available after a failed auto-submit
     */
    async retry() {
        if (!this._submissionError) {
            throw new Error('No failed submission to retry');
        }
        // Reset error state and submitted flag
        this._submissionError = null;
        this.submitted = false;
        // Reset item statuses
        for (const item of this.items) {
            if (item.status === 'failed') {
                item.status = 'pending';
                delete item.error;
            }
        }
        return this.submit();
    }
    /**
     * Add an item to the batch queue
     * Returns a placeholder that will be resolved after batch completion
     */
    add(prompt, options) {
        if (this.submitted) {
            throw new Error('Cannot add items to a submitted batch');
        }
        const item = {
            id: options?.customId || `item_${++this.idCounter}`,
            prompt,
            ...(options?.schema !== undefined && { schema: options.schema }),
            ...(options?.options !== undefined && { options: options.options }),
            ...(options?.metadata !== undefined && { metadata: options.metadata }),
            status: 'pending',
        };
        this.items.push(item);
        // Auto-submit if we hit the limit
        if (this.options.autoSubmit && this.items.length >= (this.options.maxItems || 50000)) {
            // Create a trackable promise for auto-submit
            this._autoSubmitPromise = this.submit()
                .then(() => {
                // Success - promise is resolved
            })
                .catch((error) => {
                // Store error and update item statuses
                this._submissionError = error;
                this.submitted = false; // Reset to allow retry
                // Update all items to failed status
                for (const item of this.items) {
                    if (item.status === 'pending') {
                        item.status = 'failed';
                        item.error = error.message;
                    }
                }
                // Create a synthetic failed job to store error info
                const errorWithMeta = error;
                this.job = {
                    id: `failed_${Date.now()}`,
                    provider: this.options.provider || 'openai',
                    status: 'failed',
                    totalItems: this.items.length,
                    completedItems: 0,
                    failedItems: this.items.length,
                    createdAt: new Date(),
                    // Add rate limit retry info if available
                    ...(errorWithMeta.headers?.['retry-after'] && {
                        retryAfter: parseInt(errorWithMeta.headers['retry-after'], 10),
                    }),
                };
                // Emit error event
                this.emit('error', error);
                // Log error and re-throw
                getLogger().error('Batch auto-submit failed:', error);
                throw error;
            });
        }
        return item;
    }
    /**
     * Get all items in the queue
     */
    getItems() {
        return [...this.items];
    }
    /**
     * Get the number of items in the queue
     */
    get size() {
        return this.items.length;
    }
    /**
     * Check if the batch has been submitted
     */
    get isSubmitted() {
        return this.submitted;
    }
    /**
     * Get the batch job info (after submission)
     */
    getJob() {
        return this.job;
    }
    /**
     * Submit the batch to the provider
     */
    async submit() {
        if (this.submitted) {
            throw new Error('Batch has already been submitted');
        }
        if (this.items.length === 0) {
            throw new Error('Cannot submit empty batch');
        }
        this.submitted = true;
        // Get the appropriate batch adapter
        const adapter = getBatchAdapter(this.options.provider || 'openai');
        // Submit the batch
        const result = await adapter.submit(this.items, this.options);
        this.job = result.job;
        // When completion resolves, update item statuses and check for partial failures
        result.completion.then((results) => {
            const failedResults = [];
            for (const r of results) {
                const item = this.items.find((i) => i.id === r.id);
                if (item) {
                    item.status = r.status;
                    if (r.result !== undefined)
                        item.result = r.result;
                    if (r.error !== undefined)
                        item.error = r.error;
                    if (r.status === 'failed') {
                        failedResults.push(r);
                    }
                }
            }
            // Emit partial-failure event if some items failed
            if (failedResults.length > 0 && failedResults.length < results.length) {
                this.emit('partial-failure', failedResults);
            }
            // Emit error event if there were any failures
            if (failedResults.length > 0) {
                this.emit('error', new Error(`${failedResults.length} items failed in batch`));
            }
            // Emit complete event
            this.emit('complete', results);
        });
        return result;
    }
    /**
     * Cancel the batch (if supported by provider)
     */
    async cancel() {
        if (!this.job) {
            throw new Error('Batch has not been submitted');
        }
        const adapter = getBatchAdapter(this.options.provider || 'openai');
        await adapter.cancel(this.job.id);
        this.job.status = 'cancelling';
    }
    /**
     * Get the current status of the batch
     */
    async getStatus() {
        if (!this.job) {
            throw new Error('Batch has not been submitted');
        }
        const adapter = getBatchAdapter(this.options.provider || 'openai');
        this.job = await adapter.getStatus(this.job.id);
        return this.job;
    }
    /**
     * Wait for the batch to complete and return results
     */
    async wait(pollInterval = 5000) {
        if (!this.job) {
            throw new Error('Batch has not been submitted');
        }
        const adapter = getBatchAdapter(this.options.provider || 'openai');
        return adapter.waitForCompletion(this.job.id, pollInterval);
    }
}
// Adapter registry
const adapters = {
    openai: null,
    anthropic: null,
    google: null,
    bedrock: null,
    cloudflare: null,
};
// Flex adapter registry (only OpenAI and Bedrock support flex)
const flexAdapters = {
    openai: null,
    anthropic: null,
    google: null,
    bedrock: null,
    cloudflare: null,
};
/**
 * Register a batch adapter for a provider
 *
 * Call this to register a custom batch adapter for a provider.
 * This is typically done by provider-specific packages.
 *
 * @param provider - The provider to register the adapter for
 * @param adapter - The batch adapter implementation
 *
 * @example
 * ```ts
 * import { registerBatchAdapter } from 'ai-functions'
 * import { OpenAIBatchAdapter } from './openai-adapter'
 *
 * registerBatchAdapter('openai', new OpenAIBatchAdapter())
 * ```
 */
export function registerBatchAdapter(provider, adapter) {
    adapters[provider] = adapter;
}
/**
 * Register a flex adapter for a provider
 *
 * Flex adapters provide faster-than-batch processing with the same cost savings.
 * Only OpenAI and AWS Bedrock currently support flex processing.
 *
 * @param provider - The provider to register the adapter for
 * @param adapter - The flex adapter implementation
 */
export function registerFlexAdapter(provider, adapter) {
    flexAdapters[provider] = adapter;
}
/**
 * Get the batch adapter for a provider
 *
 * @param provider - The provider to get the adapter for
 * @returns The registered batch adapter
 * @throws Error if no adapter is registered for the provider
 */
export function getBatchAdapter(provider) {
    const adapter = adapters[provider];
    if (!adapter) {
        throw new Error(`No batch adapter registered for provider: ${provider}. ` +
            `Import the adapter: import 'ai-functions/batch/${provider}'`);
    }
    return adapter;
}
/**
 * Get the flex adapter for a provider
 *
 * @param provider - The provider to get the adapter for
 * @returns The registered flex adapter
 * @throws Error if no flex adapter is registered (flex not supported by provider)
 */
export function getFlexAdapter(provider) {
    const adapter = flexAdapters[provider];
    if (!adapter) {
        throw new Error(`No flex adapter registered for provider: ${provider}. ` +
            `Flex processing is only available for OpenAI and AWS Bedrock.`);
    }
    return adapter;
}
/**
 * Check if flex processing is available for a provider
 *
 * @param provider - The provider to check
 * @returns true if flex adapter is registered, false otherwise
 */
export function hasFlexAdapter(provider) {
    return flexAdapters[provider] !== null;
}
// ============================================================================
// Tier eligibility (per-model policy)
// ============================================================================
/**
 * List the batch tiers a model is eligible for.
 *
 * Reads `ModelPolicy.batchTier` from `language-models` — this is the per-model
 * policy data, distinct from the runtime adapter registration above.
 *
 * @example
 * ```ts
 * tiersForModel('sonnet')   // ['immediate', 'batch']
 * tiersForModel('gpt-4o')   // ['immediate', 'flex', 'batch']
 * ```
 */
export function tiersForModel(alias) {
    return policyFor(alias).batchTier;
}
/**
 * Check whether a model is eligible for a given tier.
 */
export function modelSupportsTier(alias, tier) {
    return tiersForModel(alias).includes(tier);
}
// ============================================================================
// Factory Functions
// ============================================================================
/**
 * Create a new batch queue
 *
 * @example
 * ```ts
 * const batch = createBatch({ provider: 'openai' })
 * batch.add('Write a poem about cats')
 * batch.add('Write a poem about dogs')
 * const { job } = await batch.submit()
 * const results = await batch.wait()
 * ```
 */
export function createBatch(options) {
    return new BatchQueue(options);
}
/**
 * Execute operations in batch mode
 *
 * @example
 * ```ts
 * const results = await withBatch(async (batch) => {
 *   const titles = ['TypeScript', 'React', 'Next.js']
 *   return titles.map(title => batch.add(`Write a blog post about ${title}`))
 * })
 * ```
 */
export async function withBatch(fn, options) {
    const batch = createBatch(options);
    const items = await fn(batch);
    if (batch.size === 0) {
        return [];
    }
    const { completion } = await batch.submit();
    return completion;
}
// ============================================================================
// Deferred Execution Support
// ============================================================================
/**
 * Symbol to mark an AIPromise as batched/deferred
 *
 * Used internally to identify promises that should be processed via batch API.
 */
export const BATCH_MODE_SYMBOL = Symbol.for('ai-batch-mode');
/**
 * Check if we're in batch mode
 *
 * @param options - Options that may contain a batch queue
 * @returns true if a batch queue is present in options
 */
export function isBatchMode(options) {
    return !!options?.batch;
}
/**
 * Add an operation to the batch queue instead of executing immediately
 *
 * @typeParam T - The expected result type
 * @param batch - The batch queue to add to
 * @param prompt - The prompt to process
 * @param schema - Optional schema for structured output
 * @param options - Additional options including custom ID
 * @returns A BatchItem that will be resolved when the batch completes
 */
export function deferToBatch(batch, prompt, schema, options) {
    return batch.add(prompt, {
        ...(schema !== undefined && { schema }),
        ...(options !== undefined && { options }),
        ...(options?.customId !== undefined && { customId: options.customId }),
    });
}
//# sourceMappingURL=batch-queue.js.map