"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.ChatCohere = void 0;
const cohere_ai_1 = require("cohere-ai");
const messages_1 = require("@langchain/core/messages");
const chat_models_1 = require("@langchain/core/language_models/chat_models");
const outputs_1 = require("@langchain/core/outputs");
const messages_2 = require("@langchain/core/messages");
const env_1 = require("@langchain/core/utils/env");
function convertMessagesToCohereMessages(messages) {
    const getRole = (role) => {
        switch (role) {
            case "system":
                return "SYSTEM";
            case "human":
                return "USER";
            case "ai":
                return "CHATBOT";
            default:
                throw new Error(`Unknown message type: '${role}'. Accepted types: 'human', 'ai', 'system'`);
        }
    };
    const getContent = (content) => {
        if (typeof content === "string") {
            return content;
        }
        throw new Error(`ChatCohere does not support non text message content. Received: ${JSON.stringify(content, null, 2)}`);
    };
    return messages.map((message) => ({
        role: getRole(message._getType()),
        message: getContent(message.content),
    }));
}
/**
 * Integration with ChatCohere
 * @example
 * ```typescript
 * const model = new ChatCohere({
 *   apiKey: process.env.COHERE_API_KEY, // Default
 *   model: "command" // Default
 * });
 * const response = await model.invoke([
 *   new HumanMessage("How tall are the largest pengiuns?")
 * ]);
 * ```
 */
class ChatCohere extends chat_models_1.BaseChatModel {
    static lc_name() {
        return "ChatCohere";
    }
    constructor(fields) {
        super(fields ?? {});
        Object.defineProperty(this, "lc_serializable", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: true
        });
        Object.defineProperty(this, "client", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: void 0
        });
        Object.defineProperty(this, "model", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: "command"
        });
        Object.defineProperty(this, "temperature", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: 0.3
        });
        Object.defineProperty(this, "streaming", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: false
        });
        const token = fields?.apiKey ?? (0, env_1.getEnvironmentVariable)("COHERE_API_KEY");
        if (!token) {
            throw new Error("No API key provided for ChatCohere.");
        }
        this.client = new cohere_ai_1.CohereClient({
            token,
        });
        this.model = fields?.model ?? this.model;
        this.temperature = fields?.temperature ?? this.temperature;
        this.streaming = fields?.streaming ?? this.streaming;
    }
    _llmType() {
        return "cohere";
    }
    invocationParams(options) {
        const params = {
            model: this.model,
            preamble: options.preamble,
            conversationId: options.conversationId,
            promptTruncation: options.promptTruncation,
            connectors: options.connectors,
            searchQueriesOnly: options.searchQueriesOnly,
            documents: options.documents,
            temperature: options.temperature ?? this.temperature,
        };
        // Filter undefined entries
        return Object.fromEntries(Object.entries(params).filter(([, value]) => value !== undefined));
    }
    /** @ignore */
    async _generate(messages, options, runManager) {
        const tokenUsage = {};
        const params = this.invocationParams(options);
        const cohereMessages = convertMessagesToCohereMessages(messages);
        // The last message in the array is the most recent, all other messages
        // are apart of the chat history.
        const { message } = cohereMessages[cohereMessages.length - 1];
        const chatHistory = [];
        if (cohereMessages.length > 1) {
            chatHistory.push(...cohereMessages.slice(0, -1));
        }
        const input = {
            ...params,
            message,
            chatHistory,
        };
        // Handle streaming
        if (this.streaming) {
            const stream = this._streamResponseChunks(messages, options, runManager);
            const finalChunks = {};
            for await (const chunk of stream) {
                const index = chunk.generationInfo?.completion ?? 0;
                if (finalChunks[index] === undefined) {
                    finalChunks[index] = chunk;
                }
                else {
                    finalChunks[index] = finalChunks[index].concat(chunk);
                }
            }
            const generations = Object.entries(finalChunks)
                .sort(([aKey], [bKey]) => parseInt(aKey, 10) - parseInt(bKey, 10))
                .map(([_, value]) => value);
            return { generations, llmOutput: { estimatedTokenUsage: tokenUsage } };
        }
        // Not streaming, so we can just call the API once.
        const response = await this.caller.callWithOptions({ signal: options.signal }, async () => {
            let response;
            try {
                response = await this.client.chat(input);
                // eslint-disable-next-line @typescript-eslint/no-explicit-any
            }
            catch (e) {
                e.status = e.status ?? e.statusCode;
                throw e;
            }
            return response;
        });
        if ("token_count" in response) {
            const { response_tokens: completionTokens, prompt_tokens: promptTokens, total_tokens: totalTokens, } = response.token_count;
            if (completionTokens) {
                tokenUsage.completionTokens =
                    (tokenUsage.completionTokens ?? 0) + completionTokens;
            }
            if (promptTokens) {
                tokenUsage.promptTokens = (tokenUsage.promptTokens ?? 0) + promptTokens;
            }
            if (totalTokens) {
                tokenUsage.totalTokens = (tokenUsage.totalTokens ?? 0) + totalTokens;
            }
        }
        const generationInfo = { ...response };
        delete generationInfo.text;
        const generations = [
            {
                text: response.text,
                message: new messages_1.AIMessage({
                    content: response.text,
                    additional_kwargs: generationInfo,
                }),
                generationInfo,
            },
        ];
        return {
            generations,
            llmOutput: { estimatedTokenUsage: tokenUsage },
        };
    }
    async *_streamResponseChunks(messages, options, runManager) {
        const params = this.invocationParams(options);
        const cohereMessages = convertMessagesToCohereMessages(messages);
        // The last message in the array is the most recent, all other messages
        // are apart of the chat history.
        const { message } = cohereMessages[cohereMessages.length - 1];
        const chatHistory = [];
        if (cohereMessages.length > 1) {
            chatHistory.push(...cohereMessages.slice(0, -1));
        }
        const input = {
            ...params,
            message,
            chatHistory,
        };
        // All models have a built in `this.caller` property for retries
        const stream = await this.caller.call(async () => {
            let stream;
            try {
                stream = await this.client.chatStream(input);
                // eslint-disable-next-line @typescript-eslint/no-explicit-any
            }
            catch (e) {
                e.status = e.status ?? e.statusCode;
                throw e;
            }
            return stream;
        });
        for await (const chunk of stream) {
            if (chunk.eventType === "text-generation") {
                yield new outputs_1.ChatGenerationChunk({
                    text: chunk.text,
                    message: new messages_2.AIMessageChunk({ content: chunk.text }),
                });
                await runManager?.handleLLMNewToken(chunk.text);
            }
            else if (chunk.eventType !== "stream-end") {
                // Used for when the user uses their RAG/Search/other API
                // and the stream takes more actions then just text generation.
                yield new outputs_1.ChatGenerationChunk({
                    text: "",
                    message: new messages_2.AIMessageChunk({
                        content: "",
                        additional_kwargs: {
                            ...chunk,
                        },
                    }),
                    generationInfo: {
                        ...chunk,
                    },
                });
            }
        }
    }
    /** @ignore */
    _combineLLMOutput(...llmOutputs) {
        return llmOutputs.reduce((acc, llmOutput) => {
            if (llmOutput && llmOutput.estimatedTokenUsage) {
                let completionTokens = acc.estimatedTokenUsage?.completionTokens ?? 0;
                let promptTokens = acc.estimatedTokenUsage?.promptTokens ?? 0;
                let totalTokens = acc.estimatedTokenUsage?.totalTokens ?? 0;
                completionTokens +=
                    llmOutput.estimatedTokenUsage.completionTokens ?? 0;
                promptTokens += llmOutput.estimatedTokenUsage.promptTokens ?? 0;
                totalTokens += llmOutput.estimatedTokenUsage.totalTokens ?? 0;
                acc.estimatedTokenUsage = {
                    completionTokens,
                    promptTokens,
                    totalTokens,
                };
            }
            return acc;
        }, {
            estimatedTokenUsage: {
                completionTokens: 0,
                promptTokens: 0,
                totalTokens: 0,
            },
        });
    }
    get lc_secrets() {
        return {
            apiKey: "COHERE_API_KEY",
            api_key: "COHERE_API_KEY",
        };
    }
    get lc_aliases() {
        return {
            apiKey: "cohere_api_key",
            api_key: "cohere_api_key",
        };
    }
}
exports.ChatCohere = ChatCohere;
