generator-begcode
Version:
Spring Boot + Angular/React/Vue in one handy generator
70 lines (69 loc) • 2.74 kB
JavaScript
import OpenAIApi from 'openai';
import { cleanOpenAIError } from '../utils/openai.js';
import { splitArray } from './utils.js';
export const DEFAULT_ADA_CONFIG = {
model: 'text-embedding-ada-002',
maxTokensPerInput: 8191,
maxInputsPerRequest: 2048,
};
export class OpenAIEmbeddingAPI {
apiKey;
logger;
tokenizer;
modelConfig;
maxRateLimitRetries;
api;
constructor(apiKey, logger, tokenizer, baseURL, modelConfig = DEFAULT_ADA_CONFIG, maxRateLimitRetries = 3) {
this.apiKey = apiKey;
this.logger = logger;
this.tokenizer = tokenizer;
this.modelConfig = modelConfig;
this.maxRateLimitRetries = maxRateLimitRetries;
this.api = new OpenAIApi({
apiKey: this.apiKey,
dangerouslyAllowBrowser: true,
baseURL,
});
}
async createEmbeddings(input, tries) {
try {
const inputs = Array.isArray(input) ? input : [input];
this.validateInput(inputs);
const batchedInputs = splitArray(inputs, this.modelConfig.maxInputsPerRequest, this.modelConfig.maxTokensPerInput, input => this.tokenizer.encode(input).length);
const results = await Promise.all(batchedInputs.map(async (inputs) => {
const { data } = await this.api.embeddings.create({
model: this.modelConfig.model,
input: inputs,
});
return data.map(innerData => {
return {
embedding: innerData.embedding,
input: inputs[innerData.index],
};
});
}));
return results.flat();
}
catch (err) {
const error = cleanOpenAIError(err);
if (typeof error === 'object') {
const maybeOpenAiError = error;
if (maybeOpenAiError.status === 429) {
await this.logger.warning('Warning: OpenAI rate limit exceeded, sleeping for 15 seconds.');
await new Promise(resolve => setTimeout(resolve, 15000));
if (!tries || tries < this.maxRateLimitRetries) {
return this.createEmbeddings(input, tries === undefined ? 0 : ++tries);
}
}
}
throw new Error(JSON.stringify(error, null, 2));
}
}
validateInput(inputs) {
for (const input of inputs) {
if (this.tokenizer.encode(input).length > this.modelConfig.maxTokensPerInput) {
throw new Error(`Input exceeds max request tokens: ${this.modelConfig.maxTokensPerInput}`);
}
}
}
}