ai-utils.js
Version:
Build AI applications, chatbots, and agents with JavaScript and TypeScript.
327 lines (326 loc) • 12.1 kB
JavaScript
"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.OpenAITextResponseFormat = exports.OpenAITextGenerationModel = exports.calculateOpenAITextGenerationCostInMillicents = exports.isOpenAITextGenerationModel = exports.OPENAI_TEXT_GENERATION_MODELS = void 0;
const secure_json_parse_1 = __importDefault(require("secure-json-parse"));
const zod_1 = __importDefault(require("zod"));
const AbstractModel_js_1 = require("../../model-function/AbstractModel.cjs");
const AsyncQueue_js_1 = require("../../model-function/generate-text/AsyncQueue.cjs");
const parseEventSourceReadableStream_js_1 = require("../../model-function/generate-text/parseEventSourceReadableStream.cjs");
const countTokens_js_1 = require("../../model-function/tokenize-text/countTokens.cjs");
const PromptMappingTextGenerationModel_js_1 = require("../../prompt/PromptMappingTextGenerationModel.cjs");
const callWithRetryAndThrottle_js_1 = require("../../util/api/callWithRetryAndThrottle.cjs");
const postToApi_js_1 = require("../../util/api/postToApi.cjs");
const OpenAIError_js_1 = require("./OpenAIError.cjs");
const TikTokenTokenizer_js_1 = require("./TikTokenTokenizer.cjs");
/**
* @see https://platform.openai.com/docs/models/
* @see https://openai.com/pricing
*/
exports.OPENAI_TEXT_GENERATION_MODELS = {
"text-davinci-003": {
contextWindowSize: 4096,
tokenCostInMillicents: 2,
},
"text-davinci-002": {
contextWindowSize: 4096,
tokenCostInMillicents: 2,
},
"code-davinci-002": {
contextWindowSize: 8000,
tokenCostInMillicents: 2,
},
davinci: {
contextWindowSize: 2048,
tokenCostInMillicents: 2,
},
"text-curie-001": {
contextWindowSize: 2048,
tokenCostInMillicents: 0.2,
},
curie: {
contextWindowSize: 2048,
tokenCostInMillicents: 0.2,
},
"text-babbage-001": {
contextWindowSize: 2048,
tokenCostInMillicents: 0.05,
},
babbage: {
contextWindowSize: 2048,
tokenCostInMillicents: 0.05,
},
"text-ada-001": {
contextWindowSize: 2048,
tokenCostInMillicents: 0.04,
},
ada: {
contextWindowSize: 2048,
tokenCostInMillicents: 0.04,
},
};
const isOpenAITextGenerationModel = (model) => model in exports.OPENAI_TEXT_GENERATION_MODELS;
exports.isOpenAITextGenerationModel = isOpenAITextGenerationModel;
const calculateOpenAITextGenerationCostInMillicents = ({ model, response, }) => response.usage.total_tokens *
exports.OPENAI_TEXT_GENERATION_MODELS[model].tokenCostInMillicents;
exports.calculateOpenAITextGenerationCostInMillicents = calculateOpenAITextGenerationCostInMillicents;
/**
* Create a text generation model that calls the OpenAI text completion API.
*
* @see https://platform.openai.com/docs/api-reference/completions/create
*
* @example
* const model = new OpenAITextGenerationModel({
* model: "text-davinci-003",
* temperature: 0.7,
* maxTokens: 500,
* retry: retryWithExponentialBackoff({ maxTries: 5 }),
* });
*
* const { text } = await generateText(
* model,
* "Write a short story about a robot learning to love:\n\n"
* );
*/
class OpenAITextGenerationModel extends AbstractModel_js_1.AbstractModel {
constructor(settings) {
super({ settings });
Object.defineProperty(this, "provider", {
enumerable: true,
configurable: true,
writable: true,
value: "openai"
});
Object.defineProperty(this, "contextWindowSize", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
Object.defineProperty(this, "tokenizer", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
this.tokenizer = new TikTokenTokenizer_js_1.TikTokenTokenizer({ model: settings.model });
this.contextWindowSize =
exports.OPENAI_TEXT_GENERATION_MODELS[settings.model].contextWindowSize;
}
get modelName() {
return this.settings.model;
}
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 countPromptTokens(input) {
return (0, countTokens_js_1.countTokens)(this.tokenizer, input);
}
async callAPI(prompt, options) {
const { run, settings, responseFormat } = options;
const callSettings = Object.assign({
apiKey: this.apiKey,
user: this.settings.isUserIdForwardingEnabled ? run?.userId : undefined,
}, this.settings, settings, {
abortSignal: run?.abortSignal,
prompt,
responseFormat,
});
return (0, callWithRetryAndThrottle_js_1.callWithRetryAndThrottle)({
retry: callSettings.retry,
throttle: callSettings.throttle,
call: async () => callOpenAITextGenerationAPI(callSettings),
});
}
generateTextResponse(prompt, options) {
return this.callAPI(prompt, {
...options,
responseFormat: exports.OpenAITextResponseFormat.json,
});
}
extractText(response) {
return response.choices[0].text;
}
generateDeltaStreamResponse(prompt, options) {
return this.callAPI(prompt, {
...options,
responseFormat: exports.OpenAITextResponseFormat.deltaIterable,
});
}
extractTextDelta(fullDelta) {
return fullDelta[0].delta;
}
mapPrompt(promptMapping) {
return new PromptMappingTextGenerationModel_js_1.PromptMappingTextGenerationModel({
model: this.withStopTokens(promptMapping.stopTokens),
promptMapping,
});
}
withSettings(additionalSettings) {
return new OpenAITextGenerationModel(Object.assign({}, this.settings, additionalSettings));
}
get maxCompletionTokens() {
return this.settings.maxTokens;
}
withMaxCompletionTokens(maxCompletionTokens) {
return this.withSettings({ maxTokens: maxCompletionTokens });
}
withStopTokens(stopTokens) {
return this.withSettings({ stop: stopTokens });
}
}
exports.OpenAITextGenerationModel = OpenAITextGenerationModel;
const openAITextGenerationResponseSchema = zod_1.default.object({
id: zod_1.default.string(),
object: zod_1.default.literal("text_completion"),
created: zod_1.default.number(),
model: zod_1.default.string(),
choices: zod_1.default.array(zod_1.default.object({
text: zod_1.default.string(),
index: zod_1.default.number(),
logprobs: zod_1.default.nullable(zod_1.default.any()),
finish_reason: zod_1.default.string(),
})),
usage: zod_1.default.object({
prompt_tokens: zod_1.default.number(),
completion_tokens: zod_1.default.number(),
total_tokens: zod_1.default.number(),
}),
});
/**
* Call the OpenAI Text Completion API to generate a text completion for the given prompt.
*
* @see https://platform.openai.com/docs/api-reference/completions/create
*
* @example
* const response = await callOpenAITextGenerationAPI({
* apiKey: OPENAI_API_KEY,
* model: "text-davinci-003",
* prompt: "Write a short story about a robot learning to love:\n\n",
* temperature: 0.7,
* maxTokens: 500,
* });
*
* console.log(response.choices[0].text);
*/
async function callOpenAITextGenerationAPI({ baseUrl = "https://api.openai.com/v1", abortSignal, responseFormat, apiKey, model, prompt, suffix, maxTokens, temperature, topP, n, logprobs, echo, stop, presencePenalty, frequencyPenalty, bestOf, user, }) {
return (0, postToApi_js_1.postJsonToApi)({
url: `${baseUrl}/completions`,
apiKey,
body: {
stream: responseFormat.stream,
model,
prompt,
suffix,
max_tokens: maxTokens,
temperature,
top_p: topP,
n,
logprobs,
echo,
stop,
presence_penalty: presencePenalty,
frequency_penalty: frequencyPenalty,
best_of: bestOf,
user,
},
failedResponseHandler: OpenAIError_js_1.failedOpenAICallResponseHandler,
successfulResponseHandler: responseFormat.handler,
abortSignal,
});
}
exports.OpenAITextResponseFormat = {
/**
* Returns the response as a JSON object.
*/
json: {
stream: false,
handler: (0, postToApi_js_1.createJsonResponseHandler)(openAITextGenerationResponseSchema),
},
/**
* Returns an async iterable over the full deltas (all choices, including full current state at time of event)
* of the response stream.
*/
deltaIterable: {
stream: true,
handler: async ({ response }) => createOpenAITextFullDeltaIterableQueue(response.body),
},
};
const textResponseStreamEventSchema = zod_1.default.object({
choices: zod_1.default.array(zod_1.default.object({
text: zod_1.default.string(),
finish_reason: zod_1.default.enum(["stop", "length"]).nullable(),
index: zod_1.default.number(),
})),
created: zod_1.default.number(),
id: zod_1.default.string(),
model: zod_1.default.string(),
object: zod_1.default.string(),
});
async function createOpenAITextFullDeltaIterableQueue(stream) {
const queue = new AsyncQueue_js_1.AsyncQueue();
const streamDelta = [];
// process the stream asynchonously (no 'await' on purpose):
(0, parseEventSourceReadableStream_js_1.parseEventSourceReadableStream)({
stream,
callback: (event) => {
if (event.type !== "event") {
return;
}
const data = event.data;
if (data === "[DONE]") {
queue.close();
return;
}
try {
const json = secure_json_parse_1.default.parse(data);
const parseResult = textResponseStreamEventSchema.safeParse(json);
if (!parseResult.success) {
queue.push({
type: "error",
error: parseResult.error,
});
queue.close();
return;
}
const event = parseResult.data;
for (let i = 0; i < event.choices.length; i++) {
const eventChoice = event.choices[i];
const delta = eventChoice.text;
if (streamDelta[i] == null) {
streamDelta[i] = {
content: "",
isComplete: false,
delta: "",
};
}
const choice = streamDelta[i];
choice.delta = delta;
if (eventChoice.finish_reason != null) {
choice.isComplete = true;
}
choice.content += delta;
}
// Since we're mutating the choices array in an async scenario,
// we need to make a deep copy:
const streamDeltaDeepCopy = JSON.parse(JSON.stringify(streamDelta));
queue.push({
type: "delta",
fullDelta: streamDeltaDeepCopy,
});
}
catch (error) {
queue.push({ type: "error", error });
queue.close();
return;
}
},
});
return queue;
}