ai-utils.js
Version:
Build AI applications, chatbots, and agents with JavaScript and TypeScript.
137 lines (136 loc) • 4.89 kB
JavaScript
;
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.CohereTokenizer = void 0;
const zod_1 = __importDefault(require("zod"));
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");
/**
* Tokenizer for the Cohere models. It uses the Co.Tokenize and Co.Detokenize APIs.
*
* @see https://docs.cohere.com/reference/tokenize
* @see https://docs.cohere.com/reference/detokenize-1
*
* @example
* const tokenizer = new CohereTokenizer({ model: "command-nightly" });
*
* const text = "At first, Nox didn't know what to do with the pup.";
*
* const tokenCount = await countTokens(tokenizer, text);
* const tokens = await tokenizer.tokenize(text);
* const tokensAndTokenTexts = await tokenizer.tokenizeWithTexts(text);
* const reconstructedText = await tokenizer.detokenize(tokens);
*/
class CohereTokenizer {
constructor(settings) {
Object.defineProperty(this, "settings", {
enumerable: true,
configurable: true,
writable: true,
value: void 0
});
this.settings = settings;
}
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 callTokenizeAPI(text, context) {
return (0, callWithRetryAndThrottle_js_1.callWithRetryAndThrottle)({
retry: this.settings.retry,
throttle: this.settings.throttle,
call: async () => callCohereTokenizeAPI({
abortSignal: context?.abortSignal,
apiKey: this.apiKey,
text,
...this.settings,
}),
});
}
async callDeTokenizeAPI(tokens, context) {
return (0, callWithRetryAndThrottle_js_1.callWithRetryAndThrottle)({
retry: this.settings.retry,
throttle: this.settings.throttle,
call: async () => callCohereDetokenizeAPI({
abortSignal: context?.abortSignal,
apiKey: this.apiKey,
tokens,
...this.settings,
}),
});
}
async tokenize(text) {
return (await this.tokenizeWithTexts(text)).tokens;
}
async tokenizeWithTexts(text) {
const response = await this.callTokenizeAPI(text);
return {
tokens: response.tokens,
tokenTexts: response.token_strings,
};
}
async detokenize(tokens) {
const response = await this.callDeTokenizeAPI(tokens);
return response.text;
}
}
exports.CohereTokenizer = CohereTokenizer;
const cohereDetokenizationResponseSchema = zod_1.default.object({
text: zod_1.default.string(),
meta: zod_1.default.object({
api_version: zod_1.default.object({
version: zod_1.default.string(),
}),
}),
});
/**
* Call the Cohere Co.Detokenize API to detokenize a text.
*
* https://docs.cohere.com/reference/detokenize-1
*/
async function callCohereDetokenizeAPI({ baseUrl = "https://api.cohere.ai/v1", abortSignal, apiKey, model, tokens, }) {
return (0, postToApi_js_1.postJsonToApi)({
url: `${baseUrl}/detokenize`,
apiKey,
body: {
model,
tokens,
},
failedResponseHandler: CohereError_js_1.failedCohereCallResponseHandler,
successfulResponseHandler: (0, postToApi_js_1.createJsonResponseHandler)(cohereDetokenizationResponseSchema),
abortSignal,
});
}
const cohereTokenizationResponseSchema = zod_1.default.object({
tokens: zod_1.default.array(zod_1.default.number()),
token_strings: zod_1.default.array(zod_1.default.string()),
meta: zod_1.default.object({
api_version: zod_1.default.object({
version: zod_1.default.string(),
}),
}),
});
/**
* Call the Cohere Co.Tokenize API to tokenize a text.
*
* https://docs.cohere.com/reference/tokenize
*/
async function callCohereTokenizeAPI({ baseUrl = "https://api.cohere.ai/v1", abortSignal, apiKey, model, text, }) {
return (0, postToApi_js_1.postJsonToApi)({
url: `${baseUrl}/tokenize`,
apiKey,
body: {
model,
text,
},
failedResponseHandler: CohereError_js_1.failedCohereCallResponseHandler,
successfulResponseHandler: (0, postToApi_js_1.createJsonResponseHandler)(cohereTokenizationResponseSchema),
abortSignal,
});
}