UNPKG

langchain-gigachat

Version:
154 lines (153 loc) 5.87 kB
"use strict"; Object.defineProperty(exports, "__esModule", { value: true }); exports.GigaChatEmbeddings = void 0; const embeddings_1 = require("@langchain/core/embeddings"); const chunk_array_1 = require("@langchain/core/utils/chunk_array"); const gigachat_1 = require("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 }); * * ``` */ class GigaChatEmbeddings extends embeddings_1.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_1.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 = (0, chunk_array_1.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; } }); } } exports.GigaChatEmbeddings = GigaChatEmbeddings;