lac.cpp

#include "arg.h"
#include "common.h"
#include "console.h"
#include "log.h"
#include "sampling.h"
#include "llama.h"

#include "lang-cli-src/config.h"
#include "lang-cli-src/file_manager.h"
#include "lang-cli-src/model_manager.h"
#include "lang-cli-src/console_manager.h"
#include "lang-cli-src/output_parser.h"
#include "lang-cli-src/shell_executor.h"
#include "lang-cli-src/str_parser.h"

#include <iostream>
#include <string>
#include <cassert>
#include <cstdio>
#include <cstring>
#include <ctime>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>
#include <regex>
#include <queue>

#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))
#include <signal.h>
#include <unistd.h>
#elif defined (_WIN32)
#define WIN32_LEAN_AND_MEAN
#ifndef NOMINMAX
#define NOMINMAX
#endif
#include <windows.h>
#include <signal.h>
#endif

#if defined(_MSC_VER)
#pragma warning(disable: 4244 4267) // possible loss of data
#endif

static llama_context           ** g_ctx;
static llama_model             ** g_model;
static gpt_sampler             ** g_smpl;
static gpt_params               * g_params;
static std::vector<llama_token> * g_input_tokens;
static std::ostringstream       * g_output_ss;
static std::vector<llama_token> * g_output_tokens;
static bool is_interacting  = false;
static bool need_insert_eot = false;

std::string command_prompt = R"(
You are a command-line tool helper designed to assist developers in generating accurate and executable shell commands. Given a user request or problem description, generate the appropriate command to solve the issue, and explain each part of the command briefly. Always ensure the solution is correct, efficient, and follows best practices. Do not repeat yourself and provide clear and concise explanations.

### Instruction 
show current directory
### Response
```shell
pwd
```
### Explanation
- Prints the current working directory.

### Instruction
list all files in the current directory
### Response
```shell
ls
```
### Explanation 
- Lists all files and directories in the current directory.

### Instruction 
show IP address
### Response 
```shell
ip addr show
```
### Explanation 
- Displays all network interfaces and their IP addresses.

### Instruction 
update my system
### Response
```shell
sudo apt update && sudo apt upgrade -y
```
### Explanation
- `sudo apt update`: Updates the list of available packages.
- `sudo apt upgrade -y`: Upgrades all installed packages without prompting for confirmation.

### Instruction 
create a folder here named my_new_music and put a new text file named awesome_playlist.md in it
### Response
```shell
mkdir my_new_music && touch my_new_music/awesome_playlist.md
```
### Explanation
- `mkdir my_new_music`: Creates a new directory named 'my_new_music'.
- `touch my_new_music/awesome_playlist.md`: Creates a new file named 'awesome_playlist.md' inside the 'my_new_music' directory.

### Instruction 
delete the folder oldies_goldies
### Response
```shell
rm -rf oldies_goldies
```
### Explanation
- `rm -rf oldies_goldies`: Recursively and forcefully deletes the 'oldies_goldies' directory and its contents.
- **Warning**: This command permanently deletes files and directories without confirmation.

### Instruction 
mistakenly pushed large files to GitHub, please remove them from the git history
### Response
```shell
git filter-branch --index-filter 'git rm -r --cached --ignore-unmatch <file/dir>' HEAD
```
### Explanation
- `git filter-branch`: Rewrites Git history by applying filters to each commit.
- `--index-filter 'git rm -r --cached --ignore-unmatch <file/dir>'`: Removes the specified file or directory from each commit.
- `HEAD`: Applies the filter to the entire history up to the current commit.
- **Note**: Replace `<file/dir>` with the actual file or directory you want to remove.

### Instruction 
create a Python 3.11 conda environment
### Response
```shell
conda create -n "myenv" python=3.11
```
### Explanation
- `conda create -n "myenv" python=3.11`: Creates a new Conda environment named 'myenv' with Python version 3.11
)";

static void print_usage(int argc, char ** argv) {
    (void) argc;
    (void) argv;
    // LOG("If you'd like to explore more advanced model options, pass --model-help or --model-usage for additional usage instructions.");
    // LOG("\nexample usage:\n");
    // LOG("\n  text generation:     %s -m your_model.gguf -p \"I believe the meaning of life is\" -n 128\n", argv[0]);
    // LOG("\n  chat (conversation): %s -m your_model.gguf -p \"You are a helpful assistant\" -cnv\n", argv[0]);
    // LOG("\n");
}

static void write_logfile(
    const llama_context * ctx, const gpt_params & params, const llama_model * model,
    const std::vector<llama_token> & input_tokens, const std::string & output,
    const std::vector<llama_token> & output_tokens
) {
    if (params.logdir.empty()) {
        return;
    }
    
    const std::string timestamp = string_get_sortable_timestamp();

    const bool success = fs_create_directory_with_parents(params.logdir);
    if (!success) {
        // LOG_ERR("%s: failed to create logdir %s, cannot write logfile\n", __func__, params.logdir.c_str());
        return;
    }

    const std::string logfile_path = params.logdir + timestamp + ".yml";
    FILE * logfile = fopen(logfile_path.c_str(), "w");

    if (logfile == NULL) {
        // LOG_ERR("%s: failed to open logfile %s\n", __func__, logfile_path.c_str());
        return;
    }

    fprintf(logfile, "binary: main\n");
    char model_desc[128];
    llama_model_desc(model, model_desc, sizeof(model_desc));
    yaml_dump_non_result_info(logfile, params, ctx, timestamp, input_tokens, model_desc);

    fprintf(logfile, "\n");
    fprintf(logfile, "######################\n");
    fprintf(logfile, "# Generation Results #\n");
    fprintf(logfile, "######################\n");
    fprintf(logfile, "\n");

    yaml_dump_string_multiline(logfile, "output", output.c_str());
    yaml_dump_vector_int(logfile, "output_tokens", output_tokens);

    llama_perf_dump_yaml(logfile, ctx);
    fclose(logfile);
}

#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) || defined (_WIN32)
static void sigint_handler(int signo) {
    if (signo == SIGINT) {
        if (!is_interacting && g_params->interactive) {
            is_interacting  = true;
            need_insert_eot = true;
        } else {
            console::cleanup();
            LOG("Successfully exited\n");
            // LOG("\n");
            // gpt_perf_print(*g_ctx, *g_smpl);
            // write_logfile(*g_ctx, *g_params, *g_model, *g_input_tokens, g_output_ss->str(), *g_output_tokens);

            // // make sure all logs are flushed
            // LOG("Interrupted by user\n");
            // gpt_log_pause(gpt_log_main());

            _exit(130);
        }
    }
}
#endif

static std::string chat_add_and_format(struct llama_model * model, std::vector<llama_chat_msg> & chat_msgs, const std::string & role, const std::string & content) {
    llama_chat_msg new_msg{role, content};
    auto formatted = llama_chat_format_single(model, g_params->chat_template, chat_msgs, new_msg, role == "user");
    chat_msgs.push_back({role, content});
    // LOG_DBG("formatted: '%s'\n", formatted.c_str());
    return formatted;
}
// void insert_at_index(char**& argv, int& argc, const char* new_element, int index) {
//     // Create a new array with one extra space
//     char** new_argv = new char*[argc + 1];

//     // Deep copy elements up to the insertion point
//     for (int i = 0; i < index; ++i) {
//         new_argv[i] = new char[std::strlen(argv[i]) + 1];  // Allocate memory for each element
//         std::strcpy(new_argv[i], argv[i]);  // Copy each element
//     }

//     // Insert the new element (allocate memory and copy the string)
//     new_argv[index] = new char[std::strlen(new_element) + 1];
//     std::strcpy(new_argv[index], new_element);

//     // Deep copy the rest of the original array after the insertion point
//     for (int i = index; i < argc; ++i) {
//         new_argv[i + 1] = new char[std::strlen(argv[i]) + 1];  // Allocate memory for each element
//         std::strcpy(new_argv[i + 1], argv[i]);  // Copy each element
//     }

//     // Don't deallocate the original argv (as it might be system-provided)
//     // Simply replace argv with new_argv

//     // Update the original pointer and size
//     argv = new_argv;
//     ++argc;  // Increase the argument count
// }

int main(int argc, char ** argv) {
    gpt_params params;
    g_params = &params;
    model_manager mm = model_manager();
    // std::string saved_llmc_args_str = mm.get_args("llmc_args_str");
    // std::vector<std::string> saved_llmc_args = split_str(trim(saved_llmc_args_str), ' ');

    // if (argc == 1) {
    //     // no input arguments
    // } else {
    //     // // special case, only 1 argument, and it's a prompt
    //     // if (argc == 2){
    //     //     insert_at_index(argv, argc, (char*)"-p", 1);
    //     // }

    //     for (auto it = saved_llmc_args.rbegin(); it != saved_llmc_args.rend(); ++it) {
    //         insert_at_index(argv, argc, (char*)it->c_str(), 1);
    //     }
    // }
    
    if (!gpt_params_parse(argc, argv, params, LLAMA_EXAMPLE_MAIN, print_usage)) {
        return 1;
    }

    gpt_init();

    auto & sparams = params.sparams;
    // save choice to use color for later
    // (note for later: this is a slightly awkward choice)
    console::init(params.simple_io, params.use_color);
    atexit([]() { console::cleanup(); });

    if (params.logits_all) {
        // LOG_ERR("************\n");
        // LOG_ERR("%s: please use the 'perplexity' tool for perplexity calculations\n", __func__);
        // LOG_ERR("************\n\n");
        return 0;
    }

    if (params.embedding) {
        // LOG_ERR("************\n");
        // LOG_ERR("%s: please use the 'embedding' tool for embedding calculations\n", __func__);
        // LOG_ERR("************\n\n");

        return 0;
    }

    if (params.n_ctx != 0 && params.n_ctx < 8) {
        // LOG_WRN("%s: warning: minimum context size is 8, using minimum size.\n", __func__);
        params.n_ctx = 8;
    }

    if (params.rope_freq_base != 0.0) {
        // LOG_WRN("%s: warning: changing RoPE frequency base to %g.\n", __func__, params.rope_freq_base);
    }

    if (params.rope_freq_scale != 0.0) {
        // LOG_WRN("%s: warning: scaling RoPE frequency by %g.\n", __func__, params.rope_freq_scale);
    }
    // // LOG_INF("%s: llama backend init\n", __func__);

    if (params.llmc_show_args) {
        mm.show_args();
        return 0;
    }
    
    if (params.model == DEFAULT_MODEL_PATH) {
        // users do not explicitly set a model
        // std::string model_path = mm.read_model_path();
        std::string model_path = mm.get_args("model");
        if (params.llmc_setup) {
            print_centered_message("Setting up LangCommand model", PRINT_LENGTH);
            params.model = mm.set_model();
            return 0;
        }
        else if (model_path.empty() || (!file_exists(model_path) || file_is_empty(model_path))) {
            // force to reset the model
            print_centered_message("Model not found or invalid. Setting up LangCommand model", PRINT_LENGTH);
            params.model = mm.set_model();
        } else {
            params.model = model_path;
        }
    }

    // if (params.llmc_save_args) {
    //     printf("Saving args %s\n", params.llmc_args_str.c_str());
    //     mm.save_args("llmc_args_str", params.llmc_args_str);
    // }

    // llmc_reset
    if (params.prompt == "") {
        print_error("Please provide a prompt");
        return 0;
    }

    if (params.n_ctx < 4096 && params.n_ctx != 0) {
        print_error("Context size must be at least 4096");
        return 0;
    }

    params.llmc_mode = to_lower_case(params.llmc_mode);
    if (params.llmc_mode != "loop" && params.llmc_mode != "exit") {
        print_error("Invalid mode. Please choose either 'loop' or 'exit'");
        return 0;
    }

    // print_centered_message("Starting llmc", PRINT_LENGTH);


    

    params.prompt = command_prompt + "\n### Instruction\n" + params.prompt + "\n" + "### Response\n";
   

    llama_backend_init();
    llama_numa_init(params.numa);

    llama_model * model = nullptr;
    llama_context * ctx = nullptr;
    gpt_sampler * smpl = nullptr;

    std::vector<llama_chat_msg> chat_msgs;

    g_model = &model;
    g_ctx = &ctx;
    g_smpl = &smpl;

    // load the model and apply lora adapter, if any
    // // LOG_INF("%s: load the model and apply lora adapter, if any\n", __func__);
    llama_init_result llama_init = llama_init_from_gpt_params(params);

    model = llama_init.model;
    ctx = llama_init.context;

    if (model == NULL) {
        // LOG_ERR("%s: error: unable to load model\n", __func__);
        return 1;
    }

    // // LOG_INF("%s: llama threadpool init, n_threads = %d\n", __func__, (int) params.cpuparams.n_threads);

    struct ggml_threadpool_params tpp_batch =
            ggml_threadpool_params_from_cpu_params(params.cpuparams_batch);
    struct ggml_threadpool_params tpp =
            ggml_threadpool_params_from_cpu_params(params.cpuparams);

    set_process_priority(params.cpuparams.priority);

    struct ggml_threadpool * threadpool_batch = NULL;
    if (!ggml_threadpool_params_match(&tpp, &tpp_batch)) {
        threadpool_batch = ggml_threadpool_new(&tpp_batch);
        if (!threadpool_batch) {
            // LOG_ERR("%s: batch threadpool create failed : n_threads %d\n", __func__, tpp_batch.n_threads);
            return 1;
        }

        // Start the non-batch threadpool in the paused state
        tpp.paused = true;
    }

    struct ggml_threadpool * threadpool = ggml_threadpool_new(&tpp);
    if (!threadpool) {
        // LOG_ERR("%s: threadpool create failed : n_threads %d\n", __func__, tpp.n_threads);
        return 1;
    }

    llama_attach_threadpool(ctx, threadpool, threadpool_batch);

    const int n_ctx_train = llama_n_ctx_train(model);
    const int n_ctx = llama_n_ctx(ctx);

    if (n_ctx > n_ctx_train) {
        // LOG_WRN("%s: model was trained on only %d context tokens (%d specified)\n", __func__, n_ctx_train, n_ctx);
    }

    // print chat template example in conversation mode
    if (params.conversation) {
        if (params.enable_chat_template) {
            // LOG_INF("%s: chat template example:\n%s\n", __func__, llama_chat_format_example(model, params.chat_template).c_str());
        } else {
            // LOG_INF("%s: in-suffix/prefix is specified, chat template will be disabled\n", __func__);
        }
    }

    // print system information
    {
        // LOG_INF("\n");
        // LOG_INF("%s\n", gpt_params_get_system_info(params).c_str());
        // LOG_INF("\n");
    }

    std::string path_session = params.path_prompt_cache;
    std::vector<llama_token> session_tokens;

    if (!path_session.empty()) {
        // LOG_INF("%s: attempting to load saved session from '%s'\n", __func__, path_session.c_str());
        if (!file_exists(path_session)) {
            // LOG_INF("%s: session file does not exist, will create.\n", __func__);
        } else if (file_is_empty(path_session)) {
            // LOG_INF("%s: The session file is empty. A new session will be initialized.\n", __func__);
        } else {
            // The file exists and is not empty
            session_tokens.resize(n_ctx);
            size_t n_token_count_out = 0;
            if (!llama_state_load_file(ctx, path_session.c_str(), session_tokens.data(), session_tokens.capacity(), &n_token_count_out)) {
                // LOG_ERR("%s: failed to load session file '%s'\n", __func__, path_session.c_str());
                return 1;
            }
            session_tokens.resize(n_token_count_out);
            // LOG_INF("%s: loaded a session with prompt size of %d tokens\n", __func__, (int)session_tokens.size());
        }
    }

    const bool add_bos = llama_add_bos_token(model);
    if (!llama_model_has_encoder(model)) {
        GGML_ASSERT(!llama_add_eos_token(model));
    }

    // LOG_DBG("n_ctx: %d, add_bos: %d\n", n_ctx, add_bos);

    std::vector<llama_token> embd_inp;

    {
        auto prompt = (params.conversation && params.enable_chat_template && !params.prompt.empty())
            ? chat_add_and_format(model, chat_msgs, "system", params.prompt) // format the system prompt in conversation mode
            : params.prompt;
        if (params.interactive_first || !params.prompt.empty() || session_tokens.empty()) {
            // LOG_DBG("tokenize the prompt\n");
            embd_inp = ::llama_tokenize(ctx, prompt, true, true);
        } else {
            // LOG_DBG("use session tokens\n");
            embd_inp = session_tokens;
        }

        // LOG_DBG("prompt: \"%s\"\n", prompt.c_str());
        // LOG_DBG("tokens: %s\n", string_from(ctx, embd_inp).c_str());
    }

    // Should not run without any tokens
    if (embd_inp.empty()) {
        if (add_bos) {
            embd_inp.push_back(llama_token_bos(model));
            // LOG_WRN("embd_inp was considered empty and bos was added: %s\n", string_from(ctx, embd_inp).c_str());
        } else {
            // LOG_ERR("input is empty\n");
            return -1;
        }
    }

    // Tokenize negative prompt
    if ((int) embd_inp.size() > n_ctx - 4) {
        // LOG_ERR("%s: prompt is too long (%d tokens, max %d)\n", __func__, (int) embd_inp.size(), n_ctx - 4);
        return 1;
    }

    // debug message about similarity of saved session, if applicable
    size_t n_matching_session_tokens = 0;
    if (!session_tokens.empty()) {
        for (llama_token id : session_tokens) {
            if (n_matching_session_tokens >= embd_inp.size() || id != embd_inp[n_matching_session_tokens]) {
                break;
            }
            n_matching_session_tokens++;
        }
        if (params.prompt.empty() && n_matching_session_tokens == embd_inp.size()) {
            // LOG_INF("%s: using full prompt from session file\n", __func__);
        } else if (n_matching_session_tokens >= embd_inp.size()) {
            // LOG_INF("%s: session file has exact match for prompt!\n", __func__);
        } else if (n_matching_session_tokens < (embd_inp.size() / 2)) {
            // LOG_WRN("%s: session file has low similarity to prompt (%zu / %zu tokens); will mostly be reevaluated\n",
            //         __func__, n_matching_session_tokens, embd_inp.size());
        } else {
            // LOG_INF("%s: session file matches %zu / %zu tokens of prompt\n",
            //         __func__, n_matching_session_tokens, embd_inp.size());
        }

        // remove any "future" tokens that we might have inherited from the previous session
        llama_kv_cache_seq_rm(ctx, -1, n_matching_session_tokens, -1);
    }

    // LOG_DBG("recalculate the cached logits (check): embd_inp.size() %zu, n_matching_session_tokens %zu, embd_inp.size() %zu, session_tokens.size() %zu\n",
    //      embd_inp.size(), n_matching_session_tokens, embd_inp.size(), session_tokens.size());

    // if we will use the cache for the full prompt without reaching the end of the cache, force
    // reevaluation of the last token to recalculate the cached logits
    if (!embd_inp.empty() && n_matching_session_tokens == embd_inp.size() && session_tokens.size() > embd_inp.size()) {
        // LOG_DBG("recalculate the cached logits (do): session_tokens.resize( %zu )\n", embd_inp.size() - 1);

        session_tokens.resize(embd_inp.size() - 1);
    }

    // number of tokens to keep when resetting context
    if (params.n_keep < 0 || params.n_keep > (int) embd_inp.size()) {
        params.n_keep = (int)embd_inp.size();
    } else {
        params.n_keep += add_bos; // always keep the BOS token
    }

    if (params.conversation) {
        params.interactive_first = true;
    }

    // enable interactive mode if interactive start is specified
    if (params.interactive_first) {
        params.interactive = true;
    }

    if (params.verbose_prompt) {
        // LOG_INF("%s: prompt: '%s'\n", __func__, params.prompt.c_str());
        // LOG_INF("%s: number of tokens in prompt = %zu\n", __func__, embd_inp.size());
        for (int i = 0; i < (int) embd_inp.size(); i++) {
            // LOG_INF("%6d -> '%s'\n", embd_inp[i], llama_token_to_piece(ctx, embd_inp[i]).c_str());
        }

        if (params.n_keep > add_bos) {
            // LOG_INF("%s: static prompt based on n_keep: '", __func__);
            for (int i = 0; i < params.n_keep; i++) {
                // LOG_CNT("%s", llama_token_to_piece(ctx, embd_inp[i]).c_str());
            }
            // LOG_CNT("'\n");
        }
        // LOG_INF("\n");
    }

    // ctrl+C handling
    {
#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))
        struct sigaction sigint_action;
        sigint_action.sa_handler = sigint_handler;
        sigemptyset (&sigint_action.sa_mask);
        sigint_action.sa_flags = 0;
        sigaction(SIGINT, &sigint_action, NULL);
#elif defined (_WIN32)
        auto console_ctrl_handler = +[](DWORD ctrl_type) -> BOOL {
            return (ctrl_type == CTRL_C_EVENT) ? (sigint_handler(SIGINT), true) : false;
        };
        SetConsoleCtrlHandler(reinterpret_cast<PHANDLER_ROUTINE>(console_ctrl_handler), true);
#endif
    }

    if (params.interactive) {
        // LOG_INF("%s: interactive mode on.\n", __func__);

        if (!params.antiprompt.empty()) {
            for (const auto & antiprompt : params.antiprompt) {
                // LOG_INF("Reverse prompt: '%s'\n", antiprompt.c_str());
                if (params.verbose_prompt) {
                    auto tmp = ::llama_tokenize(ctx, antiprompt, false, true);
                    for (int i = 0; i < (int) tmp.size(); i++) {
                        // LOG_INF("%6d -> '%s'\n", tmp[i], llama_token_to_piece(ctx, tmp[i]).c_str());
                    }
                }
            }
        }

        if (params.input_prefix_bos) {
            // LOG_INF("Input prefix with BOS\n");
        }

        if (!params.input_prefix.empty()) {
            // LOG_INF("Input prefix: '%s'\n", params.input_prefix.c_str());
            if (params.verbose_prompt) {
                auto tmp = ::llama_tokenize(ctx, params.input_prefix, true, true);
                for (int i = 0; i < (int) tmp.size(); i++) {
                    // LOG_INF("%6d -> '%s'\n", tmp[i], llama_token_to_piece(ctx, tmp[i]).c_str());
                }
            }
        }

        if (!params.input_suffix.empty()) {
            // LOG_INF("Input suffix: '%s'\n", params.input_suffix.c_str());
            if (params.verbose_prompt) {
                auto tmp = ::llama_tokenize(ctx, params.input_suffix, false, true);
                for (int i = 0; i < (int) tmp.size(); i++) {
                    // LOG_INF("%6d -> '%s'\n", tmp[i], llama_token_to_piece(ctx, tmp[i]).c_str());
                }
            }
        }
    }

    smpl = gpt_sampler_init(model, sparams);
    if (!smpl) {
        // LOG_ERR("%s: failed to initialize sampling subsystem\n", __func__);
        return 1;
    }

    // LOG_INF("sampler seed: %u\n",     gpt_sampler_get_seed(smpl));
    // LOG_INF("sampler params: \n%s\n", sparams.print().c_str());
    // LOG_INF("sampler chain: %s\n",    gpt_sampler_print(smpl).c_str());

    // LOG_INF("generate: n_ctx = %d, n_batch = %d, n_predict = %d, n_keep = %d\n", n_ctx, params.n_batch, params.n_predict, params.n_keep);

    // group-attention state
    // number of grouped KV tokens so far (used only if params.grp_attn_n > 1)
    int ga_i = 0;

    const int ga_n = params.grp_attn_n;
    const int ga_w = params.grp_attn_w;

    if (ga_n != 1) {
        GGML_ASSERT(ga_n > 0                    && "grp_attn_n must be positive");                     // NOLINT
        GGML_ASSERT(ga_w % ga_n == 0            && "grp_attn_w must be a multiple of grp_attn_n");     // NOLINT
      //GGML_ASSERT(n_ctx_train % ga_w == 0     && "n_ctx_train must be a multiple of grp_attn_w");    // NOLINT
      //GGML_ASSERT(n_ctx >= n_ctx_train * ga_n && "n_ctx must be at least n_ctx_train * grp_attn_n"); // NOLINT
        // LOG_INF("self-extend: n_ctx_train = %d, grp_attn_n = %d, grp_attn_w = %d\n", n_ctx_train, ga_n, ga_w);
    }
    // LOG_INF("\n");

    if (params.interactive) {
        // const char * control_message;
        // if (params.multiline_input) {
        //     control_message = " - To return control to the AI, end your input with '\\'.\n"
        //                       " - To return control without starting a new line, end your input with '/'.\n";
        // } else {
        //     control_message = " - Press Return to return control to the AI.\n"
        //                       " - To return control without starting a new line, end your input with '/'.\n"
        //                       " - If you want to submit another line, end your input with '\\'.\n";
        // }
        // LOG_INF("== Running in interactive mode. ==\n");
#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) || defined (_WIN32)
        // LOG_INF(       " - Press Ctrl+C to interject at any time.\n");
#endif
        // LOG_INF(       "%s\n", control_message);

        is_interacting = params.interactive_first;
    }

    bool is_antiprompt        = false;
    bool input_echo           = true;
    bool display              = true;
    bool need_to_save_session = !path_session.empty() && n_matching_session_tokens < embd_inp.size();

    int n_past             = 0;
    int n_remain           = params.n_predict;
    int n_consumed         = 0;
    int n_session_consumed = 0;

    std::vector<int>   input_tokens;  g_input_tokens  = &input_tokens;
    std::vector<int>   output_tokens; g_output_tokens = &output_tokens;
    std::ostringstream output_ss;     g_output_ss     = &output_ss;
    std::ostringstream assistant_ss; // for storing current assistant message, used in conversation mode

    // the first thing we will do is to output the prompt, so set color accordingly
    // console::set_display(console::prompt);
    display = params.display_prompt;

    std::vector<llama_token> embd;

    // tokenized antiprompts
    std::vector<std::vector<llama_token>> antiprompt_ids;

    antiprompt_ids.reserve(params.antiprompt.size());
    for (const std::string & antiprompt : params.antiprompt) {
        antiprompt_ids.emplace_back(::llama_tokenize(ctx, antiprompt, false, true));
    }

    if (llama_model_has_encoder(model)) {
        int enc_input_size = embd_inp.size();
        llama_token * enc_input_buf = embd_inp.data();

        if (llama_encode(ctx, llama_batch_get_one(enc_input_buf, enc_input_size, 0, 0))) {
            // LOG_ERR("%s : failed to eval\n", __func__);
            return 1;
        }

        llama_token decoder_start_token_id = llama_model_decoder_start_token(model);
        if (decoder_start_token_id == -1) {
            decoder_start_token_id = llama_token_bos(model);
        }

        embd_inp.clear();
        embd_inp.push_back(decoder_start_token_id);
    }

    // Buffer system to check for "### Instruction" and subsequent "###"
    std::string output_buffer = "";
    // std::queue<std::string> unlogged_output;
    // size_t buffer_threshold = 32;
    int early_stop_pos = -1;

    while ((n_remain != 0 && !is_antiprompt) || params.interactive) {
        // predict
        if (!embd.empty()) {
            // Note: (n_ctx - 4) here is to match the logic for commandline prompt handling via
            // --prompt or --file which uses the same value.
            int max_embd_size = n_ctx - 4;

            // Ensure the input doesn't exceed the context size by truncating embd if necessary.
            if ((int) embd.size() > max_embd_size) {
                // const int skipped_tokens = (int) embd.size() - max_embd_size;
                embd.resize(max_embd_size);

                // console::set_display(console::error);
                // LOG_WRN("<<input too long: skipped %d token%s>>", skipped_tokens, skipped_tokens != 1 ? "s" : "");
                // console::set_display(console::reset);
            }

            if (ga_n == 1) {
                // infinite text generation via context shifting
                // if we run out of context:
                // - take the n_keep first tokens from the original prompt (via n_past)
                // - take half of the last (n_ctx - n_keep) tokens and recompute the logits in batches

                if (n_past + (int) embd.size() >= n_ctx) {
                    if (!params.ctx_shift){
                        // LOG_DBG("\n\n%s: context full and context shift is disabled => stopping\n", __func__);
                        break;
                    } else {
                        if (params.n_predict == -2) {
                            // LOG_DBG("\n\n%s: context full and n_predict == -%d => stopping\n", __func__, params.n_predict);
                            break;
                        }

                        const int n_left    = n_past - params.n_keep;
                        const int n_discard = n_left/2;

                        // LOG_DBG("context full, swapping: n_past = %d, n_left = %d, n_ctx = %d, n_keep = %d, n_discard = %d\n",
                        //         n_past, n_left, n_ctx, params.n_keep, n_discard);

                        llama_kv_cache_seq_rm (ctx, 0, params.n_keep            , params.n_keep + n_discard);
                        llama_kv_cache_seq_add(ctx, 0, params.n_keep + n_discard, n_past, -n_discard);

                        n_past -= n_discard;

                        // LOG_DBG("after swap: n_past = %d\n", n_past);

                        // LOG_DBG("embd: %s\n", string_from(ctx, embd).c_str());

                        // LOG_DBG("clear session path\n");
                        path_session.clear();
                    }
                }
            } else {
                // context extension via Self-Extend
                while (n_past >= ga_i + ga_w) {
                    const int ib = (ga_n*ga_i)/ga_w;
                    const int bd = (ga_w/ga_n)*(ga_n - 1);
                    const int dd = (ga_w/ga_n) - ib*bd - ga_w;

                    // LOG_DBG("\n");
                    // LOG_DBG("shift: [%6d, %6d] + %6d -> [%6d, %6d]\n", ga_i, n_past, ib*bd, ga_i + ib*bd, n_past + ib*bd);
                    // LOG_DBG("div:   [%6d, %6d] / %6d -> [%6d, %6d]\n", ga_i + ib*bd, ga_i + ib*bd + ga_w, ga_n, (ga_i + ib*bd)/ga_n, (ga_i + ib*bd + ga_w)/ga_n);
                    // LOG_DBG("shift: [%6d, %6d] + %6d -> [%6d, %6d]\n", ga_i + ib*bd + ga_w, n_past + ib*bd, dd, ga_i + ib*bd + ga_w + dd, n_past + ib*bd + dd);

                    llama_kv_cache_seq_add(ctx, 0, ga_i,                n_past,              ib*bd);
                    llama_kv_cache_seq_div(ctx, 0, ga_i + ib*bd,        ga_i + ib*bd + ga_w, ga_n);
                    llama_kv_cache_seq_add(ctx, 0, ga_i + ib*bd + ga_w, n_past + ib*bd,      dd);

                    n_past -= bd;

                    ga_i += ga_w/ga_n;

                    // LOG_DBG("\nn_past_old = %d, n_past = %d, ga_i = %d\n\n", n_past + bd, n_past, ga_i);
                }
            }

            // try to reuse a matching prefix from the loaded session instead of re-eval (via n_past)
            if (n_session_consumed < (int) session_tokens.size()) {
                size_t i = 0;
                for ( ; i < embd.size(); i++) {
                    if (embd[i] != session_tokens[n_session_consumed]) {
                        session_tokens.resize(n_session_consumed);
                        break;
                    }

                    n_past++;
                    n_session_consumed++;

                    if (n_session_consumed >= (int) session_tokens.size()) {
                        ++i;
                        break;
                    }
                }
                if (i > 0) {
                    embd.erase(embd.begin(), embd.begin() + i);
                }
            }

            for (int i = 0; i < (int) embd.size(); i += params.n_batch) {
                int n_eval = (int) embd.size() - i;
                if (n_eval > params.n_batch) {
                    n_eval = params.n_batch;
                }

                // LOG_DBG("eval: %s\n", string_from(ctx, embd).c_str());

                if (llama_decode(ctx, llama_batch_get_one(&embd[i], n_eval, n_past, 0))) {
                    // LOG_ERR("%s : failed to eval\n", __func__);
                    return 1;
                }

                n_past += n_eval;

                // LOG_DBG("n_past = %d\n", n_past);
                // Display total tokens alongside total time
                if (params.n_print > 0 && n_past % params.n_print == 0) {
                    // LOG_DBG("\n\033[31mTokens consumed so far = %d / %d \033[0m\n", n_past, n_ctx);
                }
            }

            if (!embd.empty() && !path_session.empty()) {
                session_tokens.insert(session_tokens.end(), embd.begin(), embd.end());
                n_session_consumed = session_tokens.size();
            }
        }

        embd.clear();

        if ((int) embd_inp.size() <= n_consumed && !is_interacting) {
            // optionally save the session on first sample (for faster prompt loading next time)
            if (!path_session.empty() && need_to_save_session && !params.prompt_cache_ro) {
                need_to_save_session = false;
                llama_state_save_file(ctx, path_session.c_str(), session_tokens.data(), session_tokens.size());

                // LOG_DBG("saved session to %s\n", path_session.c_str());
            }

            const llama_token id = gpt_sampler_sample(smpl, ctx, -1);

            gpt_sampler_accept(smpl, id, /* accept_grammar= */ true);

            // // LOG_DBG("last: %s\n", string_from(ctx, smpl->prev.to_vector()).c_str());

            embd.push_back(id);

            // echo this to console
            input_echo = true;

            // decrement remaining sampling budget
            --n_remain;

            // LOG_DBG("n_remain: %d\n", n_remain);
        } else {
            // some user input remains from prompt or interaction, forward it to processing
            // LOG_DBG("embd_inp.size(): %d, n_consumed: %d\n", (int) embd_inp.size(), n_consumed);
            while ((int) embd_inp.size() > n_consumed) {
                embd.push_back(embd_inp[n_consumed]);

                // push the prompt in the sampling context in order to apply repetition penalties later
                // for the prompt, we don't apply grammar rules
                gpt_sampler_accept(smpl, embd_inp[n_consumed], /* accept_grammar= */ false);

                ++n_consumed;
                if ((int) embd.size() >= params.n_batch) {
                    break;
                }
            }
        }

        // display text
        
        for (auto id : embd) {
            const std::string token_str = llama_token_to_piece(ctx, id, params.special);

            output_buffer += token_str;
            // if (output_buffer.length() <= command_prompt.length()) {
            //     // no output
            // } else {
            //     if (input_echo && display && params.llmc_show_explanation) {
            //         LOG("%s", token_str);
            //     }
            // }

            

            // Record Displayed Tokens To Log
            // Note: Generated tokens are created one by one hence this check
            if (embd.size() > 1) {
                // Incoming Requested Tokens
                input_tokens.push_back(id);
            } else {
                // Outgoing Generated Tokens
                output_tokens.push_back(id);
                output_ss << token_str;
            }

            if (output_buffer.length() >= command_prompt.length()) {
                // LOG_DBG("output_buffer: %s\n", output_buffer.c_str());
                // output_buffer.substr(command_prompt.length())
                early_stop_pos = check_early_stop(output_buffer.substr(command_prompt.length()));
                if (early_stop_pos != -1) {
                    break;
                }
            }
        }

        
        // reset color to default if there is no pending user input
        if (input_echo && (int) embd_inp.size() == n_consumed) {
            // console::set_display(console::reset);
            display = true;
        }

        // if not currently processing queued inputs;
        if ((int) embd_inp.size() <= n_consumed) {
            // check for reverse prompt in the last n_prev tokens
            if (!params.antiprompt.empty()) {
                const int n_prev = 32;
                const std::string last_output = gpt_sampler_prev_str(smpl, ctx, n_prev);

                is_antiprompt = false;
                // Check if each of the reverse prompts appears at the end of the output.
                // If we're not running interactively, the reverse prompt might be tokenized with some following characters
                // so we'll compensate for that by widening the search window a bit.
                for (std::string & antiprompt : params.antiprompt) {
                    size_t extra_padding = params.interactive ? 0 : 2;
                    size_t search_start_pos = last_output.length() > static_cast<size_t>(antiprompt.length() + extra_padding)
                        ? last_output.length() - static_cast<size_t>(antiprompt.length() + extra_padding)
                        : 0;

                    if (last_output.find(antiprompt, search_start_pos) != std::string::npos) {
                        if (params.interactive) {
                            is_interacting = true;
                        }
                        is_antiprompt = true;
                        break;
                    }
                }

                // check for reverse prompt using special tokens
                llama_token last_token = gpt_sampler_last(smpl);
                for (std::vector<llama_token> ids : antiprompt_ids) {
                    if (ids.size() == 1 && last_token == ids[0]) {
                        if (params.interactive) {
                            is_interacting = true;
                        }
                        is_antiprompt = true;
                        break;
                    }
                }

                if (is_antiprompt) {
                    // LOG_DBG("found antiprompt: %s\n", last_output.c_str());
                }
            }

            // deal with end of generation tokens in interactive mode
            if (llama_token_is_eog(model, gpt_sampler_last(smpl))) {
                // LOG_DBG("found an EOG token\n");

                if (params.interactive) {
                    if (!params.antiprompt.empty()) {
                        // tokenize and inject first reverse prompt
                        const auto first_antiprompt = ::llama_tokenize(ctx, params.antiprompt.front(), false, true);
                        embd_inp.insert(embd_inp.end(), first_antiprompt.begin(), first_antiprompt.end());
                        is_antiprompt = true;
                    }

                    if (params.enable_chat_template) {
                        chat_add_and_format(model, chat_msgs, "assistant", assistant_ss.str());
                    }
                    is_interacting = true;
                    // LOG("\n");
                }
            }

            // if current token is not EOG, we add it to current assistant message
            if (params.conversation) {
                const auto id = gpt_sampler_last(smpl);
                assistant_ss << llama_token_to_piece(ctx, id, false);
            }

            if (n_past > 0 && is_interacting) {
                // LOG_DBG("waiting for user input\n");

                if (params.conversation) {
                    LOG("\n> ");
                }

                if (params.input_prefix_bos) {
                    // LOG_DBG("adding input prefix BOS token\n");
                    embd_inp.push_back(llama_token_bos(model));
                }

                std::string buffer;
                if (!params.input_prefix.empty() && !params.conversation) {
                    // LOG_DBG("appending input prefix: '%s'\n", params.input_prefix.c_str());
                    LOG("%s", params.input_prefix.c_str());
                }

                // color user input only
                // console::set_display(console::user_input);
                display = params.display_prompt;

                std::string line;
                bool another_line = true;
                do {
                    another_line = console::readline(line, params.multiline_input);
                    buffer += line;
                } while (another_line);

                // done taking input, reset color
                // console::set_display(console::reset);
                display = true;

                // Add tokens to embd only if the input buffer is non-empty
                // Entering a empty line lets the user pass control back
                if (buffer.length() > 1) {
                    // append input suffix if any
                    if (!params.input_suffix.empty() && !params.conversation) {
                        // LOG_DBG("appending input suffix: '%s'\n", params.input_suffix.c_str());
                        LOG("%s", params.input_suffix.c_str());
                    }

                    // LOG_DBG("buffer: '%s'\n", buffer.c_str());

                    const size_t original_size = embd_inp.size();

                    if (params.escape) {
                        string_process_escapes(buffer);
                    }

                    bool format_chat = params.conversation && params.enable_chat_template;
                    std::string user_inp = format_chat
                        ? chat_add_and_format(model, chat_msgs, "user", std::move(buffer))
                        : std::move(buffer);
                    // TODO: one inconvenient of current chat template implementation is that we can't distinguish between user input and special tokens (prefix/postfix)
                    const auto line_pfx = ::llama_tokenize(ctx, params.input_prefix, false, true);
                    const auto line_inp = ::llama_tokenize(ctx, user_inp,            false, format_chat);
                    const auto line_sfx = ::llama_tokenize(ctx, params.input_suffix, false, true);

                    // LOG_DBG("input tokens: %s\n", string_from(ctx, line_inp).c_str());

                    // if user stop generation mid-way, we must add EOT to finish model's last response
                    if (need_insert_eot && format_chat) {
                        llama_token eot = llama_token_eot(model);
                        embd_inp.push_back(eot == -1 ? llama_token_eos(model) : eot);
                        need_insert_eot = false;
                    }

                    embd_inp.insert(embd_inp.end(), line_pfx.begin(), line_pfx.end());
                    embd_inp.insert(embd_inp.end(), line_inp.begin(), line_inp.end());
                    embd_inp.insert(embd_inp.end(), line_sfx.begin(), line_sfx.end());

                    for (size_t i = original_size; i < embd_inp.size(); ++i) {
                        const llama_token token = embd_inp[i];
                        output_tokens.push_back(token);
                        output_ss << llama_token_to_piece(ctx, token);
                    }

                    // reset assistant message
                    assistant_ss.str("");

                    n_remain -= line_inp.size();
                    // LOG_DBG("n_remain: %d\n", n_remain);
                } else {
                    // LOG_DBG("empty line, passing control back\n");
                }

                input_echo = false; // do not echo this again
            }

            if (n_past > 0) {
                if (is_interacting) {
                    gpt_sampler_reset(smpl);
                }
                is_interacting = false;
            }
        }

        // end of generation
        if (!embd.empty() && llama_token_is_eog(model, embd.back()) && !(params.interactive)) {
            // LOG(" [end of text]\n");
            break;
        }

        // In interactive mode, respect the maximum number of tokens and drop back to user input when reached.
        // We skip this logic when n_predict == -1 (infinite) or -2 (stop at context size).
        if (params.interactive && n_remain <= 0 && params.n_predict >= 0) {
            n_remain = params.n_predict;
            is_interacting = true;
        }
        if (early_stop_pos != -1) {
            break;
        }  
    }


    output_buffer = output_buffer.substr(command_prompt.length());
    if (early_stop_pos != -1) {
        output_buffer = output_buffer.substr(0, early_stop_pos);
        // output_buffer = output_buffer.substr(command_prompt.length());

        // while (!unlogged_output.empty()) {
        //     if (input_echo && display && params.llmc_show_explanation) {
        //         LOG("%s", unlogged_output.front().c_str());
        //         unlogged_output.pop();
        //     }
        // }
    }
    output_buffer = trim(output_buffer);
    if (params.llmc_show_explanation) {
        // printf("%s", output_buffer.c_str());
        std::string colorized_text = colorize_text(output_buffer.c_str());
        printf("%s", colorized_text.c_str());
    }
    // printf("output_buffer: %s\n", output_buffer.c_str());
    
    
    if (!path_session.empty() && params.prompt_cache_all && !params.prompt_cache_ro) {
        LOG("\n%s: saving final output to session file '%s'\n", __func__, path_session.c_str());
        llama_state_save_file(ctx, path_session.c_str(), session_tokens.data(), session_tokens.size());
    }

    // std::string output = output_ss.str();

    // std::vector<std::string> output_lines = extract_bash_blocks(output);
    printf("\n");
    std::vector<std::string> output_lines = extract_suggestions(output_buffer);
    // printf("============ Choose & Execute ============\n");
    if (output_lines.size() != 0) {
        if (params.llmc_mode == "loop") {
            while (!output_lines.empty()) {
                choose_edit_exec(output_lines);
            }
        } else if (params.llmc_mode == "exit") {
            choose_edit_exec(output_lines);
        }
    }

    gpt_sampler_free(smpl);

    llama_free(ctx);
    llama_free_model(model);

    llama_backend_free();

    ggml_threadpool_free(threadpool);
    ggml_threadpool_free(threadpool_batch);

    return 0;
}