import * as readline from "readline";
import process from "process";
import path from "path";
import chalk from "chalk";
import fs from "fs-extra";
import { chatCommandHistoryFilePath, defaultChatSystemPrompt } from "../../config.js";
import { LlamaChatWrapper } from "../../chatWrappers/LlamaChatWrapper.js";
import { GeneralChatWrapper } from "../../chatWrappers/GeneralChatWrapper.js";
import { ChatMLChatWrapper } from "../../chatWrappers/ChatMLChatWrapper.js";
import { resolveChatWrapperBasedOnModel } from "../../chatWrappers/resolveChatWrapperBasedOnModel.js";
import { FalconChatWrapper } from "../../chatWrappers/FalconChatWrapper.js";
import { getIsInDocumentationMode } from "../../state.js";
import { ReplHistory } from "../../utils/ReplHistory.js";
import withStatusLogs from "../../utils/withStatusLogs.js";
import { AlpacaChatWrapper } from "../../chatWrappers/AlpacaChatWrapper.js";
import { FunctionaryChatWrapper } from "../../chatWrappers/FunctionaryChatWrapper.js";
import { defineChatSessionFunction } from "../../evaluator/LlamaChatSession/utils/defineChatSessionFunction.js";
import { getLlama } from "../../bindings/getLlama.js";
import { LlamaGrammar } from "../../evaluator/LlamaGrammar.js";
import { LlamaChatSession } from "../../evaluator/LlamaChatSession/LlamaChatSession.js";
import { LlamaModel } from "../../evaluator/LlamaModel.js";
import { LlamaContext } from "../../evaluator/LlamaContext/LlamaContext.js";
import { LlamaJsonSchemaGrammar } from "../../evaluator/LlamaJsonSchemaGrammar.js";
import { LlamaLogLevel } from "../../bindings/types.js";
const modelWrappers = ["auto", "general", "llamaChat", "alpacaChat", "functionary", "chatML", "falconChat"];
export const ChatCommand = {
    command: "chat",
    describe: "Chat with a Llama model",
    builder(yargs) {
        const isInDocumentationMode = getIsInDocumentationMode();
        return yargs
            .option("model", {
            alias: "m",
            type: "string",
            demandOption: true,
            description: "Llama model file to use for the chat",
            group: "Required:"
        })
            .option("systemInfo", {
            alias: "i",
            type: "boolean",
            default: false,
            description: "Print llama.cpp system info",
            group: "Optional:"
        })
            .option("systemPrompt", {
            alias: "s",
            type: "string",
            default: defaultChatSystemPrompt,
            defaultDescription: " ",
            description: "System prompt to use against the model" +
                (isInDocumentationMode ? "" : (". [default value: " + defaultChatSystemPrompt.split("\n").join(" ") + "]")),
            group: "Optional:"
        })
            .option("systemPromptFile", {
            type: "string",
            description: "Path to a file to load text from and use as as the model system prompt",
            group: "Optional:"
        })
            .option("prompt", {
            type: "string",
            description: "First prompt to automatically send to the model when starting the chat",
            group: "Optional:"
        })
            .option("promptFile", {
            type: "string",
            description: "Path to a file to load text from and use as a first prompt to automatically send to the model when starting the chat",
            group: "Optional:"
        })
            .option("wrapper", {
            alias: "w",
            type: "string",
            default: "auto",
            choices: modelWrappers,
            description: "Chat wrapper to use. Use `auto` to automatically select a wrapper based on the model's BOS token",
            group: "Optional:"
        })
            .option("contextSize", {
            alias: "c",
            type: "number",
            default: 1024 * 4,
            description: "Context size to use for the model context",
            group: "Optional:"
        })
            .option("batchSize", {
            alias: "b",
            type: "number",
            description: "Batch size to use for the model context. The default value is the context size",
            group: "Optional:"
        })
            .option("grammar", {
            alias: "g",
            type: "string",
            default: "text",
            choices: ["text", "json", "list", "arithmetic", "japanese", "chess"],
            description: "Restrict the model response to a specific grammar, like JSON for example",
            group: "Optional:"
        })
            .option("jsonSchemaGrammarFile", {
            alias: ["jsgf"],
            type: "string",
            description: "File path to a JSON schema file, to restrict the model response to only generate output that conforms to the JSON schema",
            group: "Optional:"
        })
            .option("threads", {
            type: "number",
            default: 6,
            description: "Number of threads to use for the evaluation of tokens",
            group: "Optional:"
        })
            .option("temperature", {
            alias: "t",
            type: "number",
            default: 0,
            description: "Temperature is a hyperparameter that controls the randomness of the generated text. It affects the probability distribution of the model's output tokens. A higher temperature (e.g., 1.5) makes the output more random and creative, while a lower temperature (e.g., 0.5) makes the output more focused, deterministic, and conservative. The suggested temperature is 0.8, which provides a balance between randomness and determinism. At the extreme, a temperature of 0 will always pick the most likely next token, leading to identical outputs in each run. Set to `0` to disable.",
            group: "Optional:"
        })
            .option("topK", {
            alias: "k",
            type: "number",
            default: 40,
            description: "Limits the model to consider only the K most likely next tokens for sampling at each step of sequence generation. An integer number between `1` and the size of the vocabulary. Set to `0` to disable (which uses the full vocabulary). Only relevant when `temperature` is set to a value greater than 0.",
            group: "Optional:"
        })
            .option("topP", {
            alias: "p",
            type: "number",
            default: 0.95,
            description: "Dynamically selects the smallest set of tokens whose cumulative probability exceeds the threshold P, and samples the next token only from this set. A float number between `0` and `1`. Set to `1` to disable. Only relevant when `temperature` is set to a value greater than `0`.",
            group: "Optional:"
        })
            .option("gpuLayers", {
            alias: "gl",
            type: "number",
            description: "number of layers to store in VRAM",
            group: "Optional:"
        })
            .option("repeatPenalty", {
            alias: "rp",
            type: "number",
            default: 1.1,
            description: "Prevent the model from repeating the same token too much. Set to `1` to disable.",
            group: "Optional:"
        })
            .option("lastTokensRepeatPenalty", {
            alias: "rpn",
            type: "number",
            default: 64,
            description: "Number of recent tokens generated by the model to apply penalties to repetition of",
            group: "Optional:"
        })
            .option("penalizeRepeatingNewLine", {
            alias: "rpnl",
            type: "boolean",
            default: true,
            description: "Penalize new line tokens. set \"--no-penalizeRepeatingNewLine\" or \"--no-rpnl\" to disable",
            group: "Optional:"
        })
            .option("repeatFrequencyPenalty", {
            alias: "rfp",
            type: "number",
            description: "For n time a token is in the `punishTokens` array, lower its probability by `n * repeatFrequencyPenalty`. Set to a value between `0` and `1` to enable.",
            group: "Optional:"
        })
            .option("repeatPresencePenalty", {
            alias: "rpp",
            type: "number",
            description: "Lower the probability of all the tokens in the `punishTokens` array by `repeatPresencePenalty`. Set to a value between `0` and `1` to enable.",
            group: "Optional:"
        })
            .option("maxTokens", {
            alias: "mt",
            type: "number",
            default: 0,
            description: "Maximum number of tokens to generate in responses. Set to `0` to disable. Set to `-1` to set to the context size",
            group: "Optional:"
        })
            .option("noHistory", {
            alias: "nh",
            type: "boolean",
            default: false,
            description: "Don't load or save chat history",
            group: "Optional:"
        })
            .option("environmentFunctions", {
            alias: "ef",
            type: "boolean",
            default: false,
            description: "Provide access to environment functions like `getDate` and `getTime`",
            group: "Optional:"
        })
            .option("noInfoLog", {
            alias: "nl",
            type: "boolean",
            default: false,
            description: "Disable llama.cpp info logs",
            group: "Optional:"
        })
            .option("printTimings", {
            alias: "pt",
            type: "boolean",
            default: false,
            description: "Print llama.cpp timings after each response",
            group: "Optional:"
        });
    },
    async handler({ model, systemInfo, systemPrompt, systemPromptFile, prompt, promptFile, wrapper, contextSize, batchSize, grammar, jsonSchemaGrammarFile, threads, temperature, topK, topP, gpuLayers, repeatPenalty, lastTokensRepeatPenalty, penalizeRepeatingNewLine, repeatFrequencyPenalty, repeatPresencePenalty, maxTokens, noHistory, environmentFunctions, noInfoLog, printTimings }) {
        try {
            await RunChat({
                model, systemInfo, systemPrompt, systemPromptFile, prompt, promptFile, wrapper, contextSize, batchSize,
                grammar, jsonSchemaGrammarFile, threads, temperature, topK, topP, gpuLayers, lastTokensRepeatPenalty,
                repeatPenalty, penalizeRepeatingNewLine, repeatFrequencyPenalty, repeatPresencePenalty, maxTokens,
                noHistory, environmentFunctions, noInfoLog, printTimings
            });
        }
        catch (err) {
            console.error(err);
            process.exit(1);
        }
    }
};
async function RunChat({ model: modelArg, systemInfo, systemPrompt, systemPromptFile, prompt, promptFile, wrapper, contextSize, batchSize, grammar: grammarArg, jsonSchemaGrammarFile: jsonSchemaGrammarFilePath, threads, temperature, topK, topP, gpuLayers, lastTokensRepeatPenalty, repeatPenalty, penalizeRepeatingNewLine, repeatFrequencyPenalty, repeatPresencePenalty, maxTokens, noHistory, environmentFunctions, noInfoLog, printTimings }) {
    if (noInfoLog)
        console.info(`${chalk.yellow("Log level:")} warn`);
    const llama = await getLlama("lastBuild", {
        logLevel: noInfoLog
            ? LlamaLogLevel.warn
            : LlamaLogLevel.debug
    });
    const logBatchSize = batchSize != null;
    if (systemInfo)
        console.log(llama.systemInfo);
    if (systemPromptFile != null && systemPromptFile !== "") {
        if (systemPrompt != null && systemPrompt !== "" && systemPrompt !== defaultChatSystemPrompt)
            console.warn(chalk.yellow("Both `systemPrompt` and `systemPromptFile` were specified. `systemPromptFile` will be used."));
        systemPrompt = await fs.readFile(path.resolve(process.cwd(), systemPromptFile), "utf8");
    }
    if (promptFile != null && promptFile !== "") {
        if (prompt != null && prompt !== "")
            console.warn(chalk.yellow("Both `prompt` and `promptFile` were specified. `promptFile` will be used."));
        prompt = await fs.readFile(path.resolve(process.cwd(), promptFile), "utf8");
    }
    if (batchSize == null)
        batchSize = contextSize;
    else if (batchSize > contextSize) {
        console.warn(chalk.yellow("Batch size is greater than the context size. Batch size will be set to the context size."));
        batchSize = contextSize;
    }
    let initialPrompt = prompt ?? null;
    const model = await withStatusLogs({
        loading: chalk.blue("Loading model"),
        success: chalk.blue("Model loaded"),
        fail: chalk.blue("Failed to load model")
    }, async () => new LlamaModel({
        llama,
        modelPath: path.resolve(process.cwd(), modelArg),
        gpuLayers: gpuLayers != null ? gpuLayers : undefined
    }));
    const context = await withStatusLogs({
        loading: chalk.blue("Creating context"),
        success: chalk.blue("Context created"),
        fail: chalk.blue("Failed to create context")
    }, async () => new LlamaContext({
        model,
        contextSize,
        batchSize,
        threads
    }));
    const grammar = jsonSchemaGrammarFilePath != null
        ? new LlamaJsonSchemaGrammar(llama, await fs.readJson(path.resolve(process.cwd(), jsonSchemaGrammarFilePath)))
        : grammarArg !== "text"
            ? await LlamaGrammar.getFor(llama, grammarArg)
            : undefined;
    const bos = model.tokens.bosString; // bos = beginning of sequence
    const eos = model.tokens.bosString; // eos = end of sequence
    const chatWrapper = getChatWrapper(wrapper, {
        bosString: bos,
        filename: model.filename,
        typeDescription: model.typeDescription
    });
    const session = new LlamaChatSession({
        contextSequence: context.getSequence(),
        systemPrompt,
        chatWrapper: chatWrapper
    });
    await new Promise((accept) => setTimeout(accept, 0)); // wait for logs to finish printing
    if (grammarArg != "text" && jsonSchemaGrammarFilePath != null)
        console.warn(chalk.yellow("Both `grammar` and `jsonSchemaGrammarFile` were specified. `jsonSchemaGrammarFile` will be used."));
    console.info(`${chalk.yellow("Context size:")} ${context.contextSize}`);
    if (logBatchSize)
        console.info(`${chalk.yellow("Batch size:")} ${context.batchSize}`);
    console.info(`${chalk.yellow("Train context size:")} ${model.trainContextSize}`);
    console.info(`${chalk.yellow("Model type:")} ${model.typeDescription}`);
    console.info(`${chalk.yellow("BOS:")} ${bos}`);
    console.info(`${chalk.yellow("EOS:")} ${eos}`);
    console.info(`${chalk.yellow("Chat wrapper:")} ${chatWrapper.wrapperName}`);
    console.info(`${chalk.yellow("Repeat penalty:")} ${repeatPenalty} (apply to last ${lastTokensRepeatPenalty} tokens)`);
    if (repeatFrequencyPenalty != null)
        console.info(`${chalk.yellow("Repeat frequency penalty:")} ${repeatFrequencyPenalty}`);
    if (repeatPresencePenalty != null)
        console.info(`${chalk.yellow("Repeat presence penalty:")} ${repeatPresencePenalty}`);
    if (!penalizeRepeatingNewLine)
        console.info(`${chalk.yellow("Penalize repeating new line:")} disabled`);
    if (jsonSchemaGrammarFilePath != null)
        console.info(`${chalk.yellow("JSON schema grammar file:")} ${path.relative(process.cwd(), path.resolve(process.cwd(), jsonSchemaGrammarFilePath))}`);
    else if (grammarArg !== "text")
        console.info(`${chalk.yellow("Grammar:")} ${grammarArg}`);
    if (environmentFunctions && grammar != null) {
        console.warn(chalk.yellow("Environment functions are disabled since a grammar is already specified"));
        environmentFunctions = false;
    }
    // this is for ora to not interfere with readline
    await new Promise(resolve => setTimeout(resolve, 1));
    const replHistory = await ReplHistory.load(chatCommandHistoryFilePath, !noHistory);
    async function getPrompt() {
        const rl = readline.createInterface({
            input: process.stdin,
            output: process.stdout,
            history: replHistory.history.slice()
        });
        const res = await new Promise((accept) => rl.question(chalk.yellow("> "), accept));
        rl.close();
        return res;
    }
    // eslint-disable-next-line no-constant-condition
    while (true) {
        const input = initialPrompt != null
            ? initialPrompt
            : await getPrompt();
        if (initialPrompt != null) {
            console.log(chalk.green("> ") + initialPrompt);
            initialPrompt = null;
        }
        else
            await replHistory.add(input);
        if (input === ".exit")
            break;
        process.stdout.write(chalk.yellow("AI: "));
        const [startColor, endColor] = chalk.blue("MIDDLE").split("MIDDLE");
        process.stdout.write(startColor);
        await session.prompt(input, {
            grammar: grammar,
            temperature,
            topK,
            topP,
            repeatPenalty: {
                penalty: repeatPenalty,
                frequencyPenalty: repeatFrequencyPenalty != null ? repeatFrequencyPenalty : undefined,
                presencePenalty: repeatPresencePenalty != null ? repeatPresencePenalty : undefined,
                penalizeNewLine: penalizeRepeatingNewLine,
                lastTokens: lastTokensRepeatPenalty
            },
            maxTokens: maxTokens === -1
                ? context.contextSize
                : maxTokens <= 0
                    ? undefined
                    : maxTokens,
            onToken(chunk) {
                process.stdout.write(model.detokenize(chunk));
            },
            functions: (grammar == null && environmentFunctions)
                ? defaultEnvironmentFunctions
                : undefined
        });
        process.stdout.write(endColor);
        console.log();
        if (printTimings)
            await context.printTimings();
    }
}
const defaultEnvironmentFunctions = {
    getDate: defineChatSessionFunction({
        description: "Retrieve the current date",
        handler() {
            return new Date().toLocaleDateString();
        }
    }),
    getTime: defineChatSessionFunction({
        description: "Retrieve the current time",
        handler() {
            return new Date().toLocaleTimeString();
        }
    })
};
function getChatWrapper(wrapper, { bosString, filename, typeDescription }) {
    switch (wrapper) {
        case "general":
            return new GeneralChatWrapper();
        case "llamaChat":
            return new LlamaChatWrapper();
        case "alpacaChat":
            return new AlpacaChatWrapper();
        case "functionary":
            return new FunctionaryChatWrapper();
        case "chatML":
            return new ChatMLChatWrapper();
        case "falconChat":
            return new FalconChatWrapper();
        default:
    }
    if (wrapper === "auto") {
        const chatWrapper = resolveChatWrapperBasedOnModel({
            bosString,
            filename,
            typeDescription
        });
        if (chatWrapper != null)
            return new chatWrapper();
        return new GeneralChatWrapper();
    }
    void (wrapper);
    throw new Error("Unknown wrapper: " + wrapper);
}
//# sourceMappingURL=ChatCommand.js.map