"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.GoogleBaseLLM = void 0;
const manager_1 = require("@langchain/core/callbacks/manager");
const llms_1 = require("@langchain/core/language_models/llms");
const outputs_1 = require("@langchain/core/outputs");
const env_1 = require("@langchain/core/utils/env");
const connection_js_1 = require("./connection.cjs");
const common_js_1 = require("./utils/common.cjs");
const gemini_js_1 = require("./utils/gemini.cjs");
const auth_js_1 = require("./auth.cjs");
const failed_handler_js_1 = require("./utils/failed_handler.cjs");
const chat_models_js_1 = require("./chat_models.cjs");
class GoogleLLMConnection extends connection_js_1.AbstractGoogleLLMConnection {
    async formatContents(input, _parameters) {
        const parts = await this.api.messageContentToParts(input);
        const contents = [
            {
                role: "user", // Required by Vertex AI
                parts,
            },
        ];
        return contents;
    }
}
class ProxyChatGoogle extends chat_models_js_1.ChatGoogleBase {
    constructor(fields) {
        super(fields);
    }
    buildAbstractedClient(fields) {
        return fields.connection.client;
    }
}
/**
 * Integration with an LLM.
 */
class GoogleBaseLLM extends llms_1.LLM {
    // Used for tracing, replace with the same name as your class
    static lc_name() {
        return "GoogleLLM";
    }
    get lc_secrets() {
        return {
            authOptions: "GOOGLE_AUTH_OPTIONS",
        };
    }
    constructor(fields) {
        super((0, failed_handler_js_1.ensureParams)(fields));
        Object.defineProperty(this, "originalFields", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "lc_serializable", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: true
        });
        Object.defineProperty(this, "modelName", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: "gemini-pro"
        });
        Object.defineProperty(this, "model", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: "gemini-pro"
        });
        Object.defineProperty(this, "temperature", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: 0.7
        });
        Object.defineProperty(this, "maxOutputTokens", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: 1024
        });
        Object.defineProperty(this, "topP", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: 0.8
        });
        Object.defineProperty(this, "topK", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: 40
        });
        Object.defineProperty(this, "stopSequences", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: []
        });
        Object.defineProperty(this, "safetySettings", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: []
        });
        Object.defineProperty(this, "safetyHandler", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "responseMimeType", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: "text/plain"
        });
        Object.defineProperty(this, "connection", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "streamedConnection", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        this.originalFields = fields;
        (0, common_js_1.copyAndValidateModelParamsInto)(fields, this);
        this.safetyHandler =
            fields?.safetyHandler ?? new gemini_js_1.DefaultGeminiSafetyHandler();
        const client = this.buildClient(fields);
        this.buildConnection(fields ?? {}, client);
    }
    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);
        }
    }
    buildConnection(fields, client) {
        this.connection = new GoogleLLMConnection({ ...fields, ...this }, this.caller, client, false);
        this.streamedConnection = new GoogleLLMConnection({ ...fields, ...this }, this.caller, client, true);
    }
    get platform() {
        return this.connection.platform;
    }
    // Replace
    _llmType() {
        return "googlellm";
    }
    formatPrompt(prompt) {
        return prompt;
    }
    /**
     * For some given input string and options, return a string output.
     *
     * Despite the fact that `invoke` is overridden below, we still need this
     * in order to handle public APi calls to `generate()`.
     */
    async _call(prompt, options) {
        const parameters = (0, common_js_1.copyAIModelParams)(this, options);
        const result = await this.connection.request(prompt, parameters, options);
        const ret = this.connection.api.responseToString(result);
        return ret;
    }
    // Normally, you should not override this method and instead should override
    // _streamResponseChunks. We are doing so here to allow for multimodal inputs into
    // the LLM.
    async *_streamIterator(input, options) {
        // TODO: Refactor callback setup and teardown code into core
        const prompt = llms_1.BaseLLM._convertInputToPromptValue(input);
        const [runnableConfig, callOptions] = this._separateRunnableConfigFromCallOptions(options);
        const callbackManager_ = await manager_1.CallbackManager.configure(runnableConfig.callbacks, this.callbacks, runnableConfig.tags, this.tags, runnableConfig.metadata, this.metadata, { verbose: this.verbose });
        const extra = {
            options: callOptions,
            invocation_params: this?.invocationParams(callOptions),
            batch_size: 1,
        };
        const runManagers = await callbackManager_?.handleLLMStart(this.toJSON(), [prompt.toString()], undefined, undefined, extra, undefined, undefined, runnableConfig.runName);
        let generation = new outputs_1.GenerationChunk({
            text: "",
        });
        const proxyChat = this.createProxyChat();
        try {
            for await (const chunk of proxyChat._streamIterator(input, options)) {
                const stringValue = this.connection.api.chunkToString(chunk);
                const generationChunk = new outputs_1.GenerationChunk({
                    text: stringValue,
                });
                generation = generation.concat(generationChunk);
                yield stringValue;
            }
        }
        catch (err) {
            await Promise.all((runManagers ?? []).map((runManager) => runManager?.handleLLMError(err)));
            throw err;
        }
        await Promise.all((runManagers ?? []).map((runManager) => runManager?.handleLLMEnd({
            generations: [[generation]],
        })));
    }
    async predictMessages(messages, options, _callbacks) {
        const { content } = messages[0];
        const result = await this.connection.request(content, {}, options);
        const ret = this.connection.api.responseToBaseMessage(result);
        return ret;
    }
    /**
     * Internal implementation detail to allow Google LLMs to support
     * multimodal input by delegating to the chat model implementation.
     *
     * TODO: Replace with something less hacky.
     */
    createProxyChat() {
        return new ProxyChatGoogle({
            ...this.originalFields,
            connection: this.connection,
        });
    }
    // TODO: Remove the need to override this - we are doing it to
    // allow the LLM to handle multimodal types of input.
    async invoke(input, options) {
        const stream = await this._streamIterator(input, options);
        let generatedOutput = "";
        for await (const chunk of stream) {
            generatedOutput += chunk;
        }
        return generatedOutput;
    }
}
exports.GoogleBaseLLM = GoogleBaseLLM;
