"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.ChatBaiduQianfan = void 0;
const chat_models_1 = require("@langchain/core/language_models/chat_models");
const messages_1 = require("@langchain/core/messages");
const outputs_1 = require("@langchain/core/outputs");
const env_1 = require("@langchain/core/utils/env");
const qianfan_1 = require("@baiducloud/qianfan");
/**
 * Function that extracts the custom role of a generic chat message.
 * @param message Chat message from which to extract the custom role.
 * @returns The custom role of the chat message.
 */
function extractGenericMessageCustomRole(message) {
    if (message.role !== "assistant" && message.role !== "user") {
        console.warn(`Unknown message role: ${message.role}`);
    }
    return message.role;
}
/**
 * Function that converts a base message to a Qianfan message role.
 * @param message Base message to convert.
 * @returns The Qianfan message role.
 */
function messageToQianfanRole(message) {
    const type = message._getType();
    switch (type) {
        case "ai":
            return "assistant";
        case "human":
            return "user";
        case "system":
            throw new Error("System messages should not be here");
        case "function":
            throw new Error("Function messages not supported");
        case "generic": {
            if (!messages_1.ChatMessage.isInstance(message))
                throw new Error("Invalid generic chat message");
            return extractGenericMessageCustomRole(message);
        }
        default:
            throw new Error(`Unknown message type: ${type}`);
    }
}
/**
 * Wrapper around Baidu ERNIE large language models that use the Chat endpoint.
 *
 * To use you should have the `QIANFAN_AK` and `QIANFAN_SK`
 * environment variable set.
 *
 * @augments BaseLLM
 * @augments BaiduERNIEInput
 * ```
 */
class ChatBaiduQianfan extends chat_models_1.BaseChatModel {
    static lc_name() {
        return "ChatBaiduQianfan";
    }
    get callKeys() {
        return ["stop", "signal", "options"];
    }
    get lc_secrets() {
        return {
            qianfanAK: "QIANFAN_AK",
            qianfanSK: "QIANFAN_SK",
            qianfanAccessKey: "QIANFAN_ACCESS_KEY",
            qianfanSecretKey: "QIANFAN_SECRET_KEY",
        };
    }
    get lc_aliases() {
        return undefined;
    }
    constructor(fields) {
        super(fields ?? {});
        Object.defineProperty(this, "lc_serializable", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: true
        });
        Object.defineProperty(this, "streaming", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: false
        });
        Object.defineProperty(this, "prefixMessages", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "userId", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "modelName", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: "ERNIE-Bot-turbo"
        });
        Object.defineProperty(this, "model", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: "ERNIE-Bot-turbo"
        });
        Object.defineProperty(this, "temperature", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "topP", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "penaltyScore", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        Object.defineProperty(this, "client", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "qianfanAK", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "qianfanSK", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "qianfanAccessKey", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "qianfanSecretKey", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        this.modelName = fields?.model ?? fields?.modelName ?? this.model;
        this.model = this.modelName;
        if (!this.model) {
            throw new Error(`Please provide modelName`);
        }
        this.qianfanAK = fields?.qianfanAK ?? (0, env_1.getEnvironmentVariable)("QIANFAN_AK");
        this.qianfanSK = fields?.qianfanSK ?? (0, env_1.getEnvironmentVariable)("QIANFAN_SK");
        this.qianfanAccessKey =
            fields?.qianfanAccessKey ?? (0, env_1.getEnvironmentVariable)("QIANFAN_ACCESS_KEY");
        this.qianfanSecretKey =
            fields?.qianfanSecretKey ?? (0, env_1.getEnvironmentVariable)("QIANFAN_SECRET_KEY");
        // 优先使用安全认证AK/SK鉴权
        if (this.qianfanAccessKey && this.qianfanSecretKey) {
            this.client = new qianfan_1.ChatCompletion({
                QIANFAN_ACCESS_KEY: this.qianfanAccessKey,
                QIANFAN_SECRET_KEY: this.qianfanSecretKey,
            });
        }
        else if (this.qianfanAK && this.qianfanSK) {
            this.client = new qianfan_1.ChatCompletion({
                QIANFAN_AK: this.qianfanAK,
                QIANFAN_SK: this.qianfanSK,
            });
        }
        else {
            throw new Error("Please provide AK/SK");
        }
        this.streaming = fields?.streaming ?? this.streaming;
        this.prefixMessages = fields?.prefixMessages ?? this.prefixMessages;
        this.userId = fields?.userId ?? this.userId;
        this.temperature = fields?.temperature ?? this.temperature;
        this.topP = fields?.topP ?? this.topP;
        this.penaltyScore = fields?.penaltyScore ?? this.penaltyScore;
    }
    /**
     * Get the parameters used to invoke the model
     */
    invocationParams() {
        return {
            stream: this.streaming,
            user_id: this.userId,
            temperature: this.temperature,
            top_p: this.topP,
            penalty_score: this.penaltyScore,
        };
    }
    /**
     * Get the identifying parameters for the model
     */
    identifyingParams() {
        return {
            model_name: this.model,
            ...this.invocationParams(),
        };
    }
    _ensureMessages(messages) {
        return messages.map((message) => ({
            role: messageToQianfanRole(message),
            content: message.content.toString(),
        }));
    }
    /** @ignore */
    async _generate(messages, options, runManager) {
        if (this.streaming) {
            let finalChunk;
            const stream = this._streamResponseChunks(messages, options, runManager);
            for await (const chunk of stream) {
                if (finalChunk === undefined) {
                    finalChunk = chunk;
                }
                else {
                    finalChunk = finalChunk.concat(chunk);
                }
            }
            if (finalChunk === undefined) {
                throw new Error("No chunks returned from BaiduQianFan API.");
            }
            return {
                generations: [
                    {
                        text: finalChunk.text,
                        message: finalChunk.message,
                    },
                ],
                llmOutput: finalChunk.generationInfo?.usage ?? {},
            };
        }
        else {
            const params = this.invocationParams();
            const systemMessage = messages.find((message) => message._getType() === "system");
            if (systemMessage) {
                // eslint-disable-next-line no-param-reassign
                messages = messages.filter((message) => message !== systemMessage);
                params.system = systemMessage.content.toString();
            }
            const messagesMapped = this._ensureMessages(messages);
            const data = (await this.completionWithRetry({
                ...params,
                messages: messagesMapped,
            }, false));
            const tokenUsage = data.usage || {};
            const generations = [
                {
                    text: data.result || "",
                    message: new messages_1.AIMessage(data.result || ""),
                },
            ];
            return {
                generations,
                llmOutput: { tokenUsage },
            };
        }
    }
    /** @ignore */
    async completionWithRetry(request, stream) {
        const makeCompletionRequest = async () => {
            const response = await this.client.chat(request, this.model);
            if (!stream) {
                return response;
            }
            else {
                return response;
            }
        };
        return this.caller.call(makeCompletionRequest);
    }
    async *_streamResponseChunks(messages, _options, runManager) {
        const parameters = {
            ...this.invocationParams(),
            stream: true,
        };
        const systemMessage = messages.find((message) => message._getType() === "system");
        if (systemMessage) {
            // eslint-disable-next-line no-param-reassign
            messages = messages.filter((message) => message !== systemMessage);
            parameters.system = systemMessage.content.toString();
        }
        const messagesMapped = this._ensureMessages(messages);
        const stream = (await this.caller.call(async () => this.completionWithRetry({
            ...parameters,
            messages: messagesMapped,
        }, true)));
        for await (const chunk of stream) {
            const { result, is_end, id } = chunk;
            yield new outputs_1.ChatGenerationChunk({
                text: result,
                message: new messages_1.AIMessageChunk({ content: result }),
                generationInfo: is_end
                    ? {
                        is_end,
                        request_id: id,
                        usage: chunk.usage,
                    }
                    : undefined,
            });
            await runManager?.handleLLMNewToken(result);
        }
    }
    _llmType() {
        return "baiduqianfan";
    }
}
exports.ChatBaiduQianfan = ChatBaiduQianfan;
