ai-utils.js
Version:
Build AI applications, chatbots, and agents with JavaScript and TypeScript.
120 lines (119 loc) • 4.17 kB
JavaScript
import { z } from "zod";
import { AbstractModel } from "../../model-function/AbstractModel.js";
import { callWithRetryAndThrottle } from "../../util/api/callWithRetryAndThrottle.js";
import { createJsonResponseHandler, postJsonToApi, } from "../../util/api/postToApi.js";
import { failedOpenAICallResponseHandler } from "./OpenAIError.js";
/**
* @see https://openai.com/pricing
*/
const sizeToCostInMillicents = {
"1024x1024": 2000,
"512x512": 1800,
"256x256": 1600,
};
export const calculateOpenAIImageGenerationCostInMillicents = ({ settings, }) => (settings.n ?? 1) * sizeToCostInMillicents[settings.size ?? "1024x1024"];
/**
* Create an image generation model that calls the OpenAI AI image creation API.
*
* @see https://platform.openai.com/docs/api-reference/images/create
*
* @example
* const { image } = await generateImage(
* new OpenAIImageGenerationModel({ size: "512x512" }),
* "the wicked witch of the west in the style of early 19th century painting"
* );
*/
export class OpenAIImageGenerationModel extends AbstractModel {
constructor(settings) {
super({ settings });
Object.defineProperty(this, "provider", {
enumerable: true,
configurable: true,
writable: true,
value: "openai"
});
Object.defineProperty(this, "modelName", {
enumerable: true,
configurable: true,
writable: true,
value: null
});
}
get apiKey() {
const apiKey = this.settings.apiKey ?? process.env.OPENAI_API_KEY;
if (apiKey == null) {
throw new Error(`OpenAI API key is missing. Pass it as an argument to the constructor or set it as an environment variable named OPENAI_API_KEY.`);
}
return apiKey;
}
async callAPI(prompt, options) {
const run = options?.run;
const settings = options?.settings;
const responseFormat = options?.responseFormat;
const callSettings = Object.assign({
apiKey: this.apiKey,
user: this.settings.isUserIdForwardingEnabled ? run?.userId : undefined,
}, this.settings, settings, {
abortSignal: run?.abortSignal,
prompt,
responseFormat,
});
return callWithRetryAndThrottle({
retry: callSettings.retry,
throttle: callSettings.throttle,
call: async () => callOpenAIImageGenerationAPI(callSettings),
});
}
generateImageResponse(prompt, options) {
return this.callAPI(prompt, {
responseFormat: OpenAIImageGenerationResponseFormat.base64Json,
functionId: options?.functionId,
settings: options?.settings,
run: options?.run,
});
}
extractBase64Image(response) {
return response.data[0].b64_json;
}
withSettings(additionalSettings) {
return new OpenAIImageGenerationModel(Object.assign({}, this.settings, additionalSettings));
}
}
const openAIImageGenerationUrlSchema = z.object({
created: z.number(),
data: z.array(z.object({
url: z.string(),
})),
});
const openAIImageGenerationBase64JsonSchema = z.object({
created: z.number(),
data: z.array(z.object({
b64_json: z.string(),
})),
});
export const OpenAIImageGenerationResponseFormat = {
url: {
type: "url",
handler: createJsonResponseHandler(openAIImageGenerationUrlSchema),
},
base64Json: {
type: "b64_json",
handler: createJsonResponseHandler(openAIImageGenerationBase64JsonSchema),
},
};
async function callOpenAIImageGenerationAPI({ baseUrl = "https://api.openai.com/v1", abortSignal, apiKey, prompt, n, size, responseFormat, user, }) {
return postJsonToApi({
url: `${baseUrl}/images/generations`,
apiKey,
body: {
prompt,
n,
size,
response_format: responseFormat.type,
user,
},
failedResponseHandler: failedOpenAICallResponseHandler,
successfulResponseHandler: responseFormat?.handler,
abortSignal,
});
}