import { Embeddings } from "@langchain/core/embeddings";
import { chunkArray } from "@langchain/core/utils/chunk_array";
import { GigaChat } from "gigachat";
function removeEmpty(obj) {
    const newObj = {};
    for (const key in obj) {
        if (obj[key] !== undefined)
            newObj[key] = obj[key];
    }
    return newObj;
}
/**
 * Class for generating embeddings using the GigaChat API.
 * @example
 * ```typescript
 * // Embed a query using GigaChatEmbeddings to generate embeddings for a given text
 * const model = new GigaChatEmbeddings();
 * const res = await model.embedQuery(
 *   "What would be a good company name for a company that makes colorful socks?",
 * );
 * console.log({ res });
 *
 * ```
 */
export class GigaChatEmbeddings extends Embeddings {
    constructor(fields) {
        super(fields ?? {});
        Object.defineProperty(this, "prefixQuery", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: "Дано предложение, необходимо найти его парафраз \nпредложение: "
        });
        Object.defineProperty(this, "usePrefixQuery", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: false
        });
        Object.defineProperty(this, "batchSize", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: 512
        });
        Object.defineProperty(this, "stripNewLines", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: true
        });
        Object.defineProperty(this, "model", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: "Embeddings"
        });
        Object.defineProperty(this, "clientConfig", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "_client", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        this.prefixQuery = fields?.prefixQuery ?? this.prefixQuery;
        this.usePrefixQuery = fields?.usePrefixQuery ?? this.usePrefixQuery;
        this.batchSize = fields?.batchSize ?? this.batchSize;
        this.stripNewLines = fields?.stripNewLines ?? this.stripNewLines;
        this.model = fields?.model ?? this.model;
        this.clientConfig = {
            baseUrl: fields?.baseUrl,
            authUrl: fields?.authUrl,
            credentials: fields?.credentials,
            scope: fields?.scope,
            accessToken: fields?.accessToken,
            model: fields?.model,
            profanityCheck: fields?.profanityCheck,
            user: fields?.user,
            password: fields?.password,
            timeout: fields?.timeout,
            verbose: fields?.verbose,
            flags: fields?.flags,
            httpsAgent: fields?.httpsAgent,
        };
        this.clientConfig = removeEmpty(this.clientConfig);
        this._client = new GigaChat(this.clientConfig);
    }
    /**
     * Method to generate embeddings for an array of documents. Splits the
     * documents into batches and makes requests to the OpenAI API to generate
     * embeddings.
     * @param texts Array of documents to generate embeddings for.
     * @returns Promise that resolves to a 2D array of embeddings for each document.
     */
    async embedDocuments(texts) {
        const textsWithPrefix = this.usePrefixQuery
            ? texts.map((t) => this.prefixQuery + t)
            : texts;
        const batches = chunkArray(this.stripNewLines
            ? textsWithPrefix.map((t) => t.replace(/\n/g, " "))
            : textsWithPrefix, this.batchSize);
        const batchRequests = batches.map((batch) => this.embeddingWithRetry(batch));
        const batchResponses = await Promise.all(batchRequests);
        const embeddings = [];
        for (let i = 0; i < batchResponses.length; i += 1) {
            const batch = batches[i];
            const { data: batchResponse } = batchResponses[i];
            for (let j = 0; j < batch.length; j += 1) {
                embeddings.push(batchResponse[j].embedding);
            }
        }
        return embeddings;
    }
    /**
     * Method to generate an embedding for a single document. Calls the
     * embeddingWithRetry method with the document as the input.
     * @param text Document to generate an embedding for.
     * @returns Promise that resolves to an embedding for the document.
     */
    async embedQuery(text) {
        const textWithPrefix = this.usePrefixQuery ? this.prefixQuery + text : text;
        const { data } = await this.embeddingWithRetry(this.stripNewLines ? textWithPrefix.replace(/\n/g, " ") : textWithPrefix);
        return data[0].embedding;
    }
    /**
     * Private method to make a request to the GigaChat API to generate
     * embeddings. Handles the retry logic and returns the response from the
     * API.
     * @param input String or array of strings to embedding
     * @returns Promise that resolves to the response from the API.
     */
    async embeddingWithRetry(input) {
        return this.caller.call(async () => {
            try {
                const input_ = Array.isArray(input) ? input : [input];
                return await this._client.embeddings(input_, this.model);
            }
            catch (error) {
                console.error(error);
                throw error;
            }
        });
    }
}
