agsamantha/node_modules/@langchain/community/dist/embeddings/gradient_ai.js

import { Gradient } from "@gradientai/nodejs-sdk";
import { getEnvironmentVariable } from "@langchain/core/utils/env";
import { Embeddings } from "@langchain/core/embeddings";
import { chunkArray } from "@langchain/core/utils/chunk_array";
/**
 * Class for generating embeddings using the Gradient AI's API. Extends the
 * Embeddings class and implements GradientEmbeddingsParams and
 */
export class GradientEmbeddings extends Embeddings {
    constructor(fields) {
        super(fields);
        Object.defineProperty(this, "gradientAccessKey", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "workspaceId", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "batchSize", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: 128
        });
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        Object.defineProperty(this, "model", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        this.gradientAccessKey =
            fields?.gradientAccessKey ??
                getEnvironmentVariable("GRADIENT_ACCESS_TOKEN");
        this.workspaceId =
            fields?.workspaceId ?? getEnvironmentVariable("GRADIENT_WORKSPACE_ID");
        if (!this.gradientAccessKey) {
            throw new Error("Missing Gradient AI Access Token");
        }
        if (!this.workspaceId) {
            throw new Error("Missing Gradient AI Workspace ID");
        }
    }
    /**
     * Method to generate embeddings for an array of documents. Splits the
     * documents into batches and makes requests to the Gradient 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) {
        await this.setModel();
        const mappedTexts = texts.map((text) => ({ input: text }));
        const batches = chunkArray(mappedTexts, this.batchSize);
        const batchRequests = batches.map((batch) => this.caller.call(async () => this.model.generateEmbeddings({
            inputs: batch,
        })));
        const batchResponses = await Promise.all(batchRequests);
        const embeddings = [];
        for (let i = 0; i < batchResponses.length; i += 1) {
            const batch = batches[i];
            const { embeddings: 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
     * embedDocuments 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 data = await this.embedDocuments([text]);
        return data[0];
    }
    /**
     * Method to set the model to use for generating embeddings.
     * @sets the class' `model` value to that of the retrieved Embeddings Model.
     */
    async setModel() {
        if (this.model)
            return;
        const gradient = new Gradient({
            accessToken: this.gradientAccessKey,
            workspaceId: this.workspaceId,
        });
        this.model = await gradient.getEmbeddingsModel({
            slug: "bge-large",
        });
    }
}