"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.BaseGoogleEmbeddings = void 0;
const embeddings_1 = require("@langchain/core/embeddings");
const chunk_array_1 = require("@langchain/core/utils/chunk_array");
const env_1 = require("@langchain/core/utils/env");
const connection_js_1 = require("./connection.cjs");
const auth_js_1 = require("./auth.cjs");
class EmbeddingsConnection extends connection_js_1.GoogleAIConnection {
    constructor(fields, caller, client, streaming) {
        super(fields, caller, client, streaming);
        Object.defineProperty(this, "convertSystemMessageToHumanContent", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
    }
    buildUrlMethodAiStudio() {
        return "embedContent";
    }
    buildUrlMethodVertex() {
        return "predict";
    }
    async buildUrlMethod() {
        switch (this.platform) {
            case "gcp":
                return this.buildUrlMethodVertex();
            case "gai":
                return this.buildUrlMethodAiStudio();
            default:
                throw new Error(`Unknown platform when building method: ${this.platform}`);
        }
    }
    get modelPublisher() {
        // All the embedding models are currently published by "google"
        return "google";
    }
    formatDataAiStudio(input, parameters) {
        const parts = input.map((instance) => ({
            text: instance.content,
        }));
        const content = {
            parts,
        };
        const outputDimensionality = parameters?.outputDimensionality;
        const ret = {
            content,
            outputDimensionality,
        };
        // Remove undefined attributes
        let key;
        for (key in ret) {
            if (ret[key] === undefined) {
                delete ret[key];
            }
        }
        return ret;
    }
    formatDataVertex(input, parameters) {
        return {
            instances: input,
            parameters,
        };
    }
    async formatData(input, parameters) {
        switch (this.platform) {
            case "gcp":
                return this.formatDataVertex(input, parameters);
            case "gai":
                return this.formatDataAiStudio(input, parameters);
            default:
                throw new Error(`Unknown platform to format embeddings ${this.platform}`);
        }
    }
}
/**
 * Enables calls to Google APIs for generating
 * text embeddings.
 */
class BaseGoogleEmbeddings extends embeddings_1.Embeddings {
    constructor(fields) {
        super(fields);
        Object.defineProperty(this, "model", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "dimensions", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "connection", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        this.model = fields.model;
        this.dimensions = fields.dimensions ?? fields.outputDimensionality;
        this.connection = new EmbeddingsConnection({ ...fields, ...this }, this.caller, this.buildClient(fields), false);
    }
    buildApiKeyClient(apiKey) {
        return new auth_js_1.ApiKeyGoogleAuth(apiKey);
    }
    buildApiKey(fields) {
        return fields?.apiKey ?? (0, env_1.getEnvironmentVariable)("GOOGLE_API_KEY");
    }
    buildClient(fields) {
        const apiKey = this.buildApiKey(fields);
        if (apiKey) {
            return this.buildApiKeyClient(apiKey);
        }
        else {
            return this.buildAbstractedClient(fields);
        }
    }
    buildParameters() {
        const ret = {
            outputDimensionality: this.dimensions,
        };
        // Remove undefined attributes
        let key;
        for (key in ret) {
            if (ret[key] === undefined) {
                delete ret[key];
            }
        }
        return ret;
    }
    vertexResponseToValues(response) {
        const predictions = response?.data?.predictions ?? [];
        return predictions.map((prediction) => prediction.embeddings.values);
    }
    aiStudioResponseToValues(response) {
        const value = response?.data?.embedding?.values ?? [];
        return [value];
    }
    responseToValues(response) {
        switch (this.connection.platform) {
            case "gcp":
                return this.vertexResponseToValues(response);
            case "gai":
                return this.aiStudioResponseToValues(response);
            default:
                throw new Error(`Unknown response platform: ${this.connection.platform}`);
        }
    }
    /**
     * Takes an array of documents as input and returns a promise that
     * resolves to a 2D array of embeddings for each document. It splits the
     * documents into chunks and makes requests to the Google Vertex AI API to
     * generate embeddings.
     * @param documents An array of documents to be embedded.
     * @returns A promise that resolves to a 2D array of embeddings for each document.
     */
    async embedDocuments(documents) {
        // Vertex "text-" models could do up 5 documents at once,
        // but the "gemini-embedding-001" can only do 1.
        // AI Studio can only do a chunk size of 1.
        // TODO: Make this configurable
        const chunkSize = 1;
        const instanceChunks = (0, chunk_array_1.chunkArray)(documents.map((document) => ({
            content: document,
        })), chunkSize);
        const parameters = this.buildParameters();
        const options = {};
        const responses = await Promise.all(instanceChunks.map((instances) => this.connection.request(instances, parameters, options)));
        const result = responses?.map((response) => this.responseToValues(response)).flat() ??
            [];
        return result;
    }
    /**
     * Takes a document as input and returns a promise that resolves to an
     * embedding for the document. It calls the embedDocuments method with the
     * document as the input.
     * @param document A document to be embedded.
     * @returns A promise that resolves to an embedding for the document.
     */
    async embedQuery(document) {
        const data = await this.embedDocuments([document]);
        return data[0];
    }
}
exports.BaseGoogleEmbeddings = BaseGoogleEmbeddings;
