UNPKG

generator-begcode

Version:

Spring Boot + Angular/React/Vue in one handy generator

70 lines (69 loc) 2.74 kB
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}`); } } } }