import { ChatWrapper } from "../ChatWrapper.js";
import { BuiltinSpecialToken, LlamaText, SpecialToken } from "../utils/LlamaText.js";
// source: https://github.com/openai/openai-python/blob/120d225b91a8453e15240a49fb1c6794d8119326/chatml.md
export class ChatMLChatWrapper extends ChatWrapper {
    wrapperName = "ChatML";
    generateContextText(history, { availableFunctions, documentFunctionParams } = {}) {
        const historyWithFunctions = this.addAvailableFunctionsSystemMessageToHistory(history, availableFunctions, {
            documentParams: documentFunctionParams
        });
        const resultItems = [];
        let systemTexts = [];
        let userTexts = [];
        let modelTexts = [];
        let currentAggregateFocus = null;
        function flush() {
            if (systemTexts.length > 0 || userTexts.length > 0 || modelTexts.length > 0)
                resultItems.push({
                    system: systemTexts.join("\n\n"),
                    user: userTexts.join("\n\n"),
                    model: modelTexts.join("\n\n")
                });
            systemTexts = [];
            userTexts = [];
            modelTexts = [];
        }
        for (const item of historyWithFunctions) {
            if (item.type === "system") {
                if (currentAggregateFocus !== "system")
                    flush();
                currentAggregateFocus = "system";
                systemTexts.push(item.text);
            }
            else if (item.type === "user") {
                flush();
                currentAggregateFocus = null;
                userTexts.push(item.text);
            }
            else if (item.type === "model") {
                flush();
                currentAggregateFocus = null;
                modelTexts.push(this.generateModelResponseText(item.response));
            }
        }
        flush();
        const contextText = LlamaText(resultItems.map(({ system, user, model }, index) => {
            const isLastItem = index === resultItems.length - 1;
            return LlamaText([
                (system.length === 0)
                    ? LlamaText([])
                    : LlamaText([
                        new SpecialToken("<|im_start|>system\n"),
                        system,
                        new SpecialToken("<|im_end|>\n")
                    ]),
                (user.length === 0)
                    ? LlamaText([])
                    : LlamaText([
                        new SpecialToken("<|im_start|>user\n"),
                        user,
                        new SpecialToken("<|im_end|>\n")
                    ]),
                (model.length === 0 && !isLastItem)
                    ? LlamaText([])
                    : LlamaText([
                        new SpecialToken("<|im_start|>assistant\n"),
                        user,
                        isLastItem
                            ? LlamaText([])
                            : new SpecialToken("<|im_end|>\n")
                    ])
            ]);
        }));
        return {
            contextText,
            stopGenerationTriggers: [
                LlamaText(new BuiltinSpecialToken("EOS")),
                LlamaText(new SpecialToken("<|im_end|>")),
                LlamaText("<|im_end|>")
            ]
        };
    }
}
//# sourceMappingURL=ChatMLChatWrapper.js.map