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}`);
            }
        }
    }
}
