ai-functions
Version:
A powerful TypeScript library for building AI-powered applications with template literals and structured outputs
221 lines • 8.65 kB
JavaScript
import { streamObject } from 'ai';
import { openai } from '@ai-sdk/openai';
import { z } from 'zod';
import PQueue from 'p-queue';
function createQueue(options) {
if (!options.concurrency) {
return undefined;
}
return new PQueue({
concurrency: options.concurrency,
autoStart: true,
carryoverConcurrencyCount: true,
});
}
export function createListFunction(defaultOptions = {}) {
let currentPrompt;
let currentQueue;
let queueOptions;
const getQueue = (options) => {
if (!options.concurrency) {
return undefined;
}
// Create a new queue if options have changed
if (!currentQueue ||
!queueOptions?.concurrency ||
queueOptions.concurrency !== options.concurrency) {
currentQueue = createQueue(options);
queueOptions = options;
}
return currentQueue;
};
const executeRequest = async (prompt, options) => {
const modelParams = {
temperature: options.temperature,
maxTokens: options.maxTokens,
topP: options.topP,
frequencyPenalty: options.frequencyPenalty,
presencePenalty: options.presencePenalty,
stopSequences: options.stop ? Array.isArray(options.stop) ? options.stop : [options.stop] : undefined,
seed: options.seed
};
const performRequest = async () => {
try {
const model = options.model || openai(process.env.OPENAI_DEFAULT_MODEL || 'gpt-4o');
const streamOptions = {
model,
output: 'array',
schema: z.string(),
prompt: `Generate a list of items based on this prompt: ${prompt}`,
system: options.system,
...modelParams
};
const { elementStream } = streamObject(streamOptions);
const elements = [];
for await (const item of elementStream) {
elements.push(item);
}
return elements.join('\n');
}
catch (error) {
if (error instanceof Error) {
throw error;
}
throw new Error('Failed to generate list');
}
};
const queue = getQueue(options);
const result = queue
? await queue.add(performRequest)
: await performRequest();
return result;
};
const executeStreamingRequest = async function* (prompt, options) {
const modelParams = {
temperature: options.temperature,
maxTokens: options.maxTokens,
topP: options.topP,
frequencyPenalty: options.frequencyPenalty,
presencePenalty: options.presencePenalty,
stopSequences: options.stop ? Array.isArray(options.stop) ? options.stop : [options.stop] : undefined,
seed: options.seed
};
const performRequest = async function* () {
try {
const model = options.model || openai(process.env.OPENAI_DEFAULT_MODEL || 'gpt-4o');
const streamOptions = {
model,
output: 'array',
schema: z.string(),
prompt: `Generate a list of items based on this prompt: ${prompt}`,
system: options.system,
...modelParams
};
const { elementStream } = streamObject(streamOptions);
for await (const item of elementStream) {
yield item;
}
}
catch (error) {
if (error instanceof Error) {
if (error.name === 'AbortError') {
throw new Error('Stream was aborted');
}
else if (error.name === 'TimeoutError') {
throw new Error('Stream timed out');
}
else {
throw error;
}
}
throw new Error('Failed to generate list: Unknown error occurred');
}
};
const queue = getQueue(options);
if (queue) {
const operation = performRequest();
// Queue the operation based on its type
if (operation[Symbol.asyncIterator]) {
// For AsyncGenerator, queue each value
for await (const value of operation) {
yield await queue.add(async () => value);
}
}
else {
// For Promise, queue the entire operation
yield await queue.add(async () => {
const result = await operation;
return result;
});
}
}
else {
yield* performRequest();
}
};
const templateFn = async (prompt, options = defaultOptions) => {
currentPrompt = prompt;
const mergedOptions = { ...defaultOptions, ...options };
try {
return await executeRequest(prompt, mergedOptions);
}
catch (error) {
if (error instanceof Error) {
throw error;
}
throw new Error('Failed to generate list');
}
};
const asyncIterator = async function* (prompt, options = defaultOptions) {
currentPrompt = prompt;
const mergedOptions = { ...defaultOptions, ...options };
try {
for await (const item of executeStreamingRequest(prompt, mergedOptions)) {
yield item;
}
}
catch (error) {
if (error instanceof Error) {
throw error;
}
throw new Error('Failed to generate list');
}
};
function createTemplateResult(prompt, options) {
const promise = templateFn(prompt, options);
const result = Object.assign((opts) => templateFn(prompt, { ...options, ...opts }), {
[Symbol.asyncIterator]: () => asyncIterator(prompt, options),
call: (opts) => templateFn(prompt, { ...options, ...opts }),
then: promise.then.bind(promise),
catch: promise.catch.bind(promise),
finally: promise.finally.bind(promise),
});
return result;
}
function createBaseFunction() {
function fn(stringsOrOptions, ...values) {
if (!stringsOrOptions) {
return createTemplateResult('', defaultOptions);
}
if (Array.isArray(stringsOrOptions)) {
const strings = stringsOrOptions;
if (strings.length - 1 !== values.length) {
throw new Error('Template literal slots must match provided values');
}
const lastValue = values[values.length - 1];
const options = typeof lastValue === 'object' && !Array.isArray(lastValue) && lastValue !== null
? { ...defaultOptions, ...lastValue }
: defaultOptions;
const actualValues = typeof lastValue === 'object' && !Array.isArray(lastValue) && lastValue !== null
? values.slice(0, -1)
: values;
const prompt = strings.reduce((acc, str, i) => acc + str + (actualValues[i] || ''), '');
currentPrompt = prompt;
return createTemplateResult(prompt, options);
}
const options = { ...defaultOptions, ...stringsOrOptions };
return createTemplateResult('', options);
}
// Add the asyncIterator to the base function
Object.defineProperty(fn, Symbol.asyncIterator, {
value: () => asyncIterator(currentPrompt || '', defaultOptions),
enumerable: false,
configurable: true
});
// Add the withOptions method
Object.defineProperty(fn, 'withOptions', {
value: (opts) => templateFn(currentPrompt || '', { ...defaultOptions, ...opts }),
enumerable: false,
configurable: true
});
// Add the queue property
Object.defineProperty(fn, 'queue', {
get: () => currentQueue,
enumerable: false,
configurable: true
});
return fn;
}
return createBaseFunction();
}
//# sourceMappingURL=list.js.map