ai-utils.js
Version:
Build AI applications, chatbots, and agents with JavaScript and TypeScript.
172 lines (171 loc) • 6.71 kB
JavaScript
;
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.OpenAITextEmbeddingModel = exports.calculateOpenAIEmbeddingCostInMillicents = exports.isOpenAIEmbeddingModel = exports.OPENAI_TEXT_EMBEDDING_MODELS = void 0;
const zod_1 = __importDefault(require("zod"));
const AbstractModel_js_1 = require("../../model-function/AbstractModel.cjs");
const countTokens_js_1 = require("../../model-function/tokenize-text/countTokens.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");
exports.OPENAI_TEXT_EMBEDDING_MODELS = {
"text-embedding-ada-002": {
contextWindowSize: 8192,
embeddingDimensions: 1536,
tokenCostInMillicents: 0.01,
},
};
const isOpenAIEmbeddingModel = (model) => model in exports.OPENAI_TEXT_EMBEDDING_MODELS;
exports.isOpenAIEmbeddingModel = isOpenAIEmbeddingModel;
const calculateOpenAIEmbeddingCostInMillicents = ({ model, responses, }) => {
let amountInMilliseconds = 0;
for (const response of responses) {
amountInMilliseconds +=
response.usage.total_tokens *
exports.OPENAI_TEXT_EMBEDDING_MODELS[model].tokenCostInMillicents;
}
return amountInMilliseconds;
};
exports.calculateOpenAIEmbeddingCostInMillicents = calculateOpenAIEmbeddingCostInMillicents;
/**
* Create a text embedding model that calls the OpenAI embedding API.
*
* @see https://platform.openai.com/docs/api-reference/embeddings
*
* @example
* const { embeddings } = await embedTexts(
* new OpenAITextEmbeddingModel({ model: "text-embedding-ada-002" }),
* [
* "At first, Nox didn't know what to do with the pup.",
* "He keenly observed and absorbed everything around him, from the birds in the sky to the trees in the forest.",
* ]
* );
*/
class OpenAITextEmbeddingModel extends AbstractModel_js_1.AbstractModel {
constructor(settings) {
super({ settings });
Object.defineProperty(this, "provider", {
enumerable: true,
configurable: true,
writable: true,
value: "openai"
});
Object.defineProperty(this, "maxTextsPerCall", {
enumerable: true,
configurable: true,
writable: true,
value: 1
});
Object.defineProperty(this, "embeddingDimensions", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
Object.defineProperty(this, "tokenizer", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
Object.defineProperty(this, "contextWindowSize", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
this.tokenizer = new TikTokenTokenizer_js_1.TikTokenTokenizer({ model: this.modelName });
this.contextWindowSize =
exports.OPENAI_TEXT_EMBEDDING_MODELS[this.modelName].contextWindowSize;
this.embeddingDimensions =
exports.OPENAI_TEXT_EMBEDDING_MODELS[this.modelName].embeddingDimensions;
}
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 countTokens(input) {
return (0, countTokens_js_1.countTokens)(this.tokenizer, input);
}
async callAPI(text, options) {
const run = options?.run;
const settings = options?.settings;
const callSettings = Object.assign({
apiKey: this.apiKey,
user: this.settings.isUserIdForwardingEnabled ? run?.userId : undefined,
}, this.settings, settings, {
abortSignal: run?.abortSignal,
input: text,
});
return (0, callWithRetryAndThrottle_js_1.callWithRetryAndThrottle)({
retry: this.settings.retry,
throttle: this.settings.throttle,
call: async () => callOpenAITextEmbeddingAPI(callSettings),
});
}
generateEmbeddingResponse(texts, options) {
if (texts.length > this.maxTextsPerCall) {
throw new Error(`The OpenAI embedding API only supports ${this.maxTextsPerCall} texts per API call.`);
}
return this.callAPI(texts[0], options);
}
extractEmbeddings(response) {
return [response.data[0].embedding];
}
withSettings(additionalSettings) {
return new OpenAITextEmbeddingModel(Object.assign({}, this.settings, additionalSettings));
}
}
exports.OpenAITextEmbeddingModel = OpenAITextEmbeddingModel;
const openAITextEmbeddingResponseSchema = zod_1.default.object({
object: zod_1.default.literal("list"),
data: zod_1.default
.array(zod_1.default.object({
object: zod_1.default.literal("embedding"),
embedding: zod_1.default.array(zod_1.default.number()),
index: zod_1.default.number(),
}))
.length(1),
model: zod_1.default.string(),
usage: zod_1.default.object({
prompt_tokens: zod_1.default.number(),
total_tokens: zod_1.default.number(),
}),
});
/**
* Call the OpenAI Embedding API to generate an embedding for the given input.
*
* @see https://platform.openai.com/docs/api-reference/embeddings
*
* @example
* const response = await callOpenAITextEmbeddingAPI({
* apiKey: OPENAI_API_KEY,
* model: "text-embedding-ada-002",
* input: "At first, Nox didn't know what to do with the pup.",
* });
*
* console.log(response.data[0].embedding);
*/
async function callOpenAITextEmbeddingAPI({ baseUrl = "https://api.openai.com/v1", abortSignal, apiKey, model, input, user, }) {
return (0, postToApi_js_1.postJsonToApi)({
url: `${baseUrl}/embeddings`,
apiKey,
body: {
model,
input,
user,
},
failedResponseHandler: OpenAIError_js_1.failedOpenAICallResponseHandler,
successfulResponseHandler: (0, postToApi_js_1.createJsonResponseHandler)(openAITextEmbeddingResponseSchema),
abortSignal,
});
}