ai-functions
Version:
A powerful TypeScript library for building AI-powered applications with template literals and structured outputs
287 lines • 12 kB
JavaScript
import { generateText, generateObject, Output } from 'ai';
import { openai } from '@ai-sdk/openai';
import { createOpenAICompatible } from '@ai-sdk/openai-compatible';
import { Response } from 'undici';
import { z } from 'zod';
import PQueue from 'p-queue';
import { createRequestHandler } from '../utils/request-handler';
import { StreamProgressTracker } from '../utils/stream-progress';
import { AIRequestError } from '../errors';
// Add this at the top of the file, before any other code
function isTemplateStringsArray(value) {
return Array.isArray(value) && 'raw' in value;
}
export function isStreamingResult(result) {
return (result !== null &&
typeof result === 'object' &&
'experimental_stream' in result &&
typeof result.experimental_stream === 'object' &&
Symbol.asyncIterator in result.experimental_stream);
}
export function isJsonResult(result) {
return 'object' in result && 'toJsonResponse' in result;
}
// PLACEHOLDER: response creation functions
export function createAIFunction(schema) {
const fn = async (args, options = {}) => {
if (!args) {
return { schema };
}
const requestHandler = createRequestHandler({ requestHandling: options.requestHandling });
const result = await requestHandler.execute(async () => {
return generateText({
model: options.model || getProvider()('gpt-4o-mini'),
prompt: options.prompt || '',
maxRetries: 2,
experimental_output: Output.object({ schema: schema }),
system: options.system,
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,
});
});
if (!result.experimental_output) {
throw new Error('No output received from model');
}
return result.experimental_output;
};
fn.schema = schema;
return fn;
}
// PLACEHOLDER: createTemplateFunction implementation
function getProvider() {
const gateway = process.env.AI_GATEWAY;
const apiKey = process.env.OPENAI_API_KEY;
if (!apiKey) {
throw new Error('OPENAI_API_KEY environment variable is required');
}
return gateway
? createOpenAICompatible({
baseURL: gateway,
name: 'openai',
headers: {
'Authorization': `Bearer ${apiKey}`,
'Content-Type': 'application/json'
}
})
: openai;
}
export function createJsonResponse(result) {
return result.toJsonResponse();
}
export function createStreamResponse(result) {
return new Response(result.experimental_stream, {
headers: { 'Content-Type': 'text/plain' },
});
}
export function createTextResponse(result) {
return new Response(result.text, {
headers: { 'Content-Type': 'text/plain' },
});
}
export function createTemplateFunction(defaultOptions = {}) {
let currentPrompt;
let queue;
const initQueue = (options) => {
if (!queue && options.concurrency) {
queue = new PQueue({ concurrency: options.concurrency });
}
return queue;
};
const templateFn = async (prompt, options = defaultOptions) => {
currentPrompt = prompt;
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,
system: options.system,
};
const executeRequest = async () => {
try {
if (options.outputFormat === 'json') {
try {
const model = options.model || openai(process.env.OPENAI_DEFAULT_MODEL || 'gpt-4o', { structuredOutputs: true });
if (options.schema) {
const schema = options.schema instanceof z.ZodType
? options.schema
: z.object(Object.fromEntries(Object.entries(options.schema).map(([key, type]) => [
key,
type === 'string' ? z.string() :
type === 'number' ? z.number() :
type === 'boolean' ? z.boolean() :
type === 'array' ? z.array(z.unknown()) :
type === 'object' ? z.record(z.string(), z.unknown()) :
z.unknown()
])));
const result = await generateObject({
model,
schema,
prompt,
...modelParams,
});
return JSON.stringify(result.object);
}
else {
const result = await generateObject({
model,
output: 'no-schema',
prompt,
...modelParams,
});
return JSON.stringify(result.object);
}
}
catch (error) {
if (error instanceof Error) {
throw new AIRequestError(`Invalid JSON format: ${error.message}`, error, false);
}
throw new AIRequestError('Invalid JSON format', error, false);
}
}
const result = (await generateText({
model: options.model || openai('gpt-4o-mini'),
prompt,
...modelParams,
...options,
}));
if (!result) {
throw new AIRequestError('No result received from model', undefined, false);
}
return result.text;
}
catch (error) {
if (error instanceof Error) {
throw error;
}
throw new AIRequestError('Failed to generate text', error, false);
}
};
try {
const currentQueue = initQueue(options);
if (!currentQueue) {
return executeRequest();
}
return currentQueue.add(executeRequest);
}
catch (error) {
if (error instanceof Error) {
throw error;
}
throw new AIRequestError('Failed to generate text', error, false);
}
};
const asyncIterator = async function* () {
if (!currentPrompt) {
return;
}
const options = defaultOptions;
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,
system: options.system,
};
const progressTracker = options.streaming ?
new StreamProgressTracker(options.streaming) : undefined;
const executeRequest = async () => {
const result = (await generateText({
model: options.model || openai('gpt-4o-mini'),
prompt: currentPrompt,
maxRetries: 2,
...modelParams,
...options,
}));
if (!result) {
throw new Error('No result received from model');
}
return result;
};
try {
const currentQueue = initQueue(options);
const result = currentQueue
? await currentQueue.add(executeRequest)
: await executeRequest();
if (isStreamingResult(result)) {
for await (const chunk of result.experimental_stream) {
if (progressTracker) {
progressTracker.onChunk(chunk);
}
yield chunk;
}
progressTracker?.onComplete();
}
else if (result) {
if (progressTracker) {
progressTracker.onChunk(result.text);
progressTracker.onComplete();
}
yield result.text;
}
}
catch (error) {
const errorMessage = error instanceof Error ? error.message : 'Failed to generate text';
progressTracker?.onChunk(errorMessage);
progressTracker?.onComplete();
console.error('Error in asyncIterator:', error);
yield errorMessage;
}
};
const createAsyncIterablePromise = (promise, prompt, options = defaultOptions) => {
const asyncIterable = {
[Symbol.asyncIterator]: asyncIterator
};
const callablePromise = Object.assign(async (opts) => {
if (!opts)
return promise;
return templateFn(prompt, { ...options, ...opts });
}, {
then: (onfulfilled, onrejected) => promise.then(onfulfilled, onrejected),
catch: (onrejected) => promise.catch(onrejected),
finally: (onfinally) => promise.finally(onfinally),
[Symbol.asyncIterator]: asyncIterable[Symbol.asyncIterator]
});
return callablePromise;
};
const baseFn = Object.assign((stringsOrOptions, ...values) => {
if (!stringsOrOptions) {
return createAsyncIterablePromise(templateFn('', defaultOptions), '', defaultOptions);
}
if (isTemplateStringsArray(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] || ''), '');
return createAsyncIterablePromise(templateFn(prompt, options), prompt, options);
}
return createAsyncIterablePromise(templateFn('', { ...defaultOptions, ...stringsOrOptions }), '', { ...defaultOptions, ...stringsOrOptions });
}, {
withOptions: (opts = {}) => {
const prompt = opts.prompt || currentPrompt || '';
return templateFn(prompt, { ...defaultOptions, ...opts });
},
[Symbol.asyncIterator]: asyncIterator,
queue
});
return baseFn;
}
//# sourceMappingURL=index.js.map