ai-utils.js
Version:
Build AI applications, chatbots, and agents with JavaScript and TypeScript.
173 lines (172 loc) • 6.29 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.CohereTextEmbeddingModel = exports.COHERE_TEXT_EMBEDDING_MODELS = void 0;
const zod_1 = __importDefault(require("zod"));
const AbstractModel_js_1 = require("../../model-function/AbstractModel.cjs");
const callWithRetryAndThrottle_js_1 = require("../../util/api/callWithRetryAndThrottle.cjs");
const postToApi_js_1 = require("../../util/api/postToApi.cjs");
const CohereError_js_1 = require("./CohereError.cjs");
const CohereTokenizer_js_1 = require("./CohereTokenizer.cjs");
exports.COHERE_TEXT_EMBEDDING_MODELS = {
"embed-english-light-v2.0": {
contextWindowSize: 4096,
embeddingDimensions: 1024,
},
"embed-english-v2.0": {
contextWindowSize: 4096,
embeddingDimensions: 4096,
},
"embed-multilingual-v2.0": {
contextWindowSize: 4096,
embeddingDimensions: 768,
},
};
/**
* Create a text embedding model that calls the Cohere Co.Embed API.
*
* @see https://docs.cohere.com/reference/embed
*
* @example
* const { embeddings } = await embedTexts(
* new CohereTextEmbeddingModel({ model: "embed-english-light-v2.0" }),
* [
* "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 CohereTextEmbeddingModel extends AbstractModel_js_1.AbstractModel {
constructor(settings) {
super({ settings });
Object.defineProperty(this, "provider", {
enumerable: true,
configurable: true,
writable: true,
value: "cohere"
});
Object.defineProperty(this, "maxTextsPerCall", {
enumerable: true,
configurable: true,
writable: true,
value: 96
});
Object.defineProperty(this, "embeddingDimensions", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
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.contextWindowSize =
exports.COHERE_TEXT_EMBEDDING_MODELS[this.modelName].contextWindowSize;
this.tokenizer = new CohereTokenizer_js_1.CohereTokenizer({
baseUrl: this.settings.baseUrl,
apiKey: this.settings.apiKey,
model: this.settings.model,
retry: this.settings.tokenizerSettings?.retry,
throttle: this.settings.tokenizerSettings?.throttle,
});
this.embeddingDimensions =
exports.COHERE_TEXT_EMBEDDING_MODELS[this.modelName].embeddingDimensions;
}
get modelName() {
return this.settings.model;
}
async tokenize(text) {
return this.tokenizer.tokenize(text);
}
async tokenizeWithTexts(text) {
return this.tokenizer.tokenizeWithTexts(text);
}
async detokenize(tokens) {
return this.tokenizer.detokenize(tokens);
}
get apiKey() {
const apiKey = this.settings.apiKey ?? process.env.COHERE_API_KEY;
if (apiKey == null) {
throw new Error("No Cohere API key provided. Pass an API key to the constructor or set the COHERE_API_KEY environment variable.");
}
return apiKey;
}
async callAPI(texts, options) {
if (texts.length > this.maxTextsPerCall) {
throw new Error(`The Cohere embedding API only supports ${this.maxTextsPerCall} texts per API call.`);
}
const run = options?.run;
const settings = options?.settings;
const callSettings = Object.assign({
apiKey: this.apiKey,
}, this.settings, settings, {
abortSignal: run?.abortSignal,
texts,
});
return (0, callWithRetryAndThrottle_js_1.callWithRetryAndThrottle)({
retry: this.settings.retry,
throttle: this.settings.throttle,
call: async () => callCohereEmbeddingAPI(callSettings),
});
}
generateEmbeddingResponse(texts, options) {
return this.callAPI(texts, options);
}
extractEmbeddings(response) {
return response.embeddings;
}
withSettings(additionalSettings) {
return new CohereTextEmbeddingModel(Object.assign({}, this.settings, additionalSettings));
}
}
exports.CohereTextEmbeddingModel = CohereTextEmbeddingModel;
const cohereTextEmbeddingResponseSchema = zod_1.default.object({
id: zod_1.default.string(),
texts: zod_1.default.array(zod_1.default.string()),
embeddings: zod_1.default.array(zod_1.default.array(zod_1.default.number())),
meta: zod_1.default.object({
api_version: zod_1.default.object({
version: zod_1.default.string(),
}),
}),
});
/**
* Call the Cohere Co.Embed API to generate an embedding for the given input.
*
* @see https://docs.cohere.com/reference/embed
*
* @example
* const response = await callCohereEmbeddingAPI({
* apiKey: COHERE_API_KEY,
* model: "embed-english-light-v2.0",
* texts: [
* "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.",
* ],
* });
*/
async function callCohereEmbeddingAPI({ baseUrl = "https://api.cohere.ai/v1", abortSignal, apiKey, model, texts, truncate, }) {
return (0, postToApi_js_1.postJsonToApi)({
url: `${baseUrl}/embed`,
apiKey,
body: {
model,
texts,
truncate,
},
failedResponseHandler: CohereError_js_1.failedCohereCallResponseHandler,
successfulResponseHandler: (0, postToApi_js_1.createJsonResponseHandler)(cohereTextEmbeddingResponseSchema),
abortSignal,
});
}