echodict/llama.cpp

1#include "chat.h"2#include "common.h"3#include "arg.h"4#include "console.h"5#include "fit.h"6// #include "log.h"7 8#include "server-common.h"9#include "server-context.h"10#include "server-task.h"11 12#include <array>13#include <atomic>14#include <algorithm>15#include <filesystem>16#include <fstream>17#include <thread>18#include <signal.h>19 20#if defined(_WIN32)21#define WIN32_LEAN_AND_MEAN22#ifndef NOMINMAX23#   define NOMINMAX24#endif25#include <windows.h>26#endif27 28const char * LLAMA_ASCII_LOGO = R"(29▄▄ ▄▄30██ ██31██ ██  ▀▀█▄ ███▄███▄  ▀▀█▄    ▄████ ████▄ ████▄32██ ██ ▄█▀██ ██ ██ ██ ▄█▀██    ██    ██ ██ ██ ██33██ ██ ▀█▄██ ██ ██ ██ ▀█▄██ ██ ▀████ ████▀ ████▀34                                    ██    ██35                                    ▀▀    ▀▀36)";37 38static std::atomic<bool> g_is_interrupted = false;39static bool should_stop() {40    return g_is_interrupted.load();41}42 43#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) || defined (_WIN32)44static void signal_handler(int) {45    if (g_is_interrupted.load()) {46        // second Ctrl+C - exit immediately47        // make sure to clear colors before exiting (not using LOG or console.cpp here to avoid deadlock)48        fprintf(stdout, "\033[0m\n");49        fflush(stdout);50        std::exit(130);51    }52    g_is_interrupted.store(true);53}54#endif55 56struct cli_context {57    server_context ctx_server;58    json messages = json::array();59    std::vector<raw_buffer> input_files;60    task_params defaults;61    bool verbose_prompt;62    int reasoning_budget = -1;63    std::string reasoning_budget_message;64 65    // thread for showing "loading" animation66    std::atomic<bool> loading_show;67 68    cli_context(const common_params & params) {69        defaults.sampling    = params.sampling;70        defaults.speculative = params.speculative;71        defaults.n_keep      = params.n_keep;72        defaults.n_predict   = params.n_predict;73        defaults.antiprompt  = params.antiprompt;74 75        defaults.stream = true; // make sure we always use streaming mode76        defaults.timings_per_token = true; // in order to get timings even when we cancel mid-way77        // defaults.return_progress = true; // TODO: show progress78 79        verbose_prompt = params.verbose_prompt;80        reasoning_budget = params.reasoning_budget;81        reasoning_budget_message = params.reasoning_budget_message;82    }83 84    std::string generate_completion(result_timings & out_timings) {85        server_response_reader rd = ctx_server.get_response_reader();86        auto chat_params = format_chat();87        {88            // TODO: reduce some copies here in the future89            server_task task = server_task(SERVER_TASK_TYPE_COMPLETION);90            task.id         = rd.get_new_id();91            task.index      = 0;92            task.params     = defaults;           // copy93            task.cli_prompt = chat_params.prompt; // copy94            task.cli_files  = input_files;        // copy95            task.cli        = true;96 97            // chat template settings98            task.params.chat_parser_params = common_chat_parser_params(chat_params);99            task.params.chat_parser_params.reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK;100            if (!chat_params.parser.empty()) {101                task.params.chat_parser_params.parser.load(chat_params.parser);102            }103 104            // reasoning budget sampler105            if (!chat_params.thinking_end_tag.empty()) {106                const llama_vocab * vocab = llama_model_get_vocab(107                    llama_get_model(ctx_server.get_llama_context()));108 109                task.params.sampling.reasoning_budget_tokens = reasoning_budget;110                task.params.sampling.generation_prompt = chat_params.generation_prompt;111 112                if (!chat_params.thinking_start_tag.empty()) {113                    task.params.sampling.reasoning_budget_start =114                        common_tokenize(vocab, chat_params.thinking_start_tag, false, true);115                }116                task.params.sampling.reasoning_budget_end =117                    common_tokenize(vocab, chat_params.thinking_end_tag, false, true);118                task.params.sampling.reasoning_budget_forced =119                    common_tokenize(vocab, reasoning_budget_message + chat_params.thinking_end_tag, false, true);120            }121 122            rd.post_task({std::move(task)});123        }124 125        if (verbose_prompt) {126            console::set_display(DISPLAY_TYPE_PROMPT);127            console::log("%s\n\n", chat_params.prompt.c_str());128            console::set_display(DISPLAY_TYPE_RESET);129        }130 131        // wait for first result132        console::spinner::start();133        server_task_result_ptr result = rd.next(should_stop);134 135        console::spinner::stop();136        std::string curr_content;137        bool is_thinking = false;138 139        while (result) {140            if (should_stop()) {141                break;142            }143            if (result->is_error()) {144                json err_data = result->to_json();145                if (err_data.contains("message")) {146                    console::error("Error: %s\n", err_data["message"].get<std::string>().c_str());147                } else {148                    console::error("Error: %s\n", err_data.dump().c_str());149                }150                return curr_content;151            }152            auto res_partial = dynamic_cast<server_task_result_cmpl_partial *>(result.get());153            if (res_partial) {154                out_timings = std::move(res_partial->timings);155                for (const auto & diff : res_partial->oaicompat_msg_diffs) {156                    if (!diff.content_delta.empty()) {157                        if (is_thinking) {158                            console::log("\n[End thinking]\n\n");159                            console::set_display(DISPLAY_TYPE_RESET);160                            is_thinking = false;161                        }162                        curr_content += diff.content_delta;163                        console::log("%s", diff.content_delta.c_str());164                        console::flush();165                    }166                    if (!diff.reasoning_content_delta.empty()) {167                        console::set_display(DISPLAY_TYPE_REASONING);168                        if (!is_thinking) {169                            console::log("[Start thinking]\n");170                        }171                        is_thinking = true;172                        console::log("%s", diff.reasoning_content_delta.c_str());173                        console::flush();174                    }175                }176            }177            auto res_final = dynamic_cast<server_task_result_cmpl_final *>(result.get());178            if (res_final) {179                out_timings = std::move(res_final->timings);180                break;181            }182            result = rd.next(should_stop);183        }184        g_is_interrupted.store(false);185        // server_response_reader automatically cancels pending tasks upon destruction186        return curr_content;187    }188 189    // TODO: support remote files in the future (http, https, etc)190    std::string load_input_file(const std::string & fname, bool is_media) {191        std::ifstream file(fname, std::ios::binary);192        if (!file) {193            return "";194        }195        if (is_media) {196            raw_buffer buf;197            buf.assign((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());198            input_files.push_back(std::move(buf));199            return get_media_marker();200        } else {201            std::string content((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());202            return content;203        }204    }205 206    common_chat_params format_chat() {207        auto meta = ctx_server.get_meta();208        auto & chat_params = meta.chat_params;209 210        common_chat_templates_inputs inputs;211        inputs.messages              = common_chat_msgs_parse_oaicompat(messages);212        inputs.tools                 = {}; // TODO213        inputs.tool_choice           = COMMON_CHAT_TOOL_CHOICE_NONE;214        inputs.json_schema           = ""; // TODO215        inputs.grammar               = ""; // TODO216        inputs.use_jinja             = chat_params.use_jinja;217        inputs.parallel_tool_calls   = false;218        inputs.add_generation_prompt = true;219        inputs.reasoning_format      = COMMON_REASONING_FORMAT_DEEPSEEK;220        inputs.force_pure_content    = chat_params.force_pure_content;221        inputs.enable_thinking       = chat_params.enable_thinking ? common_chat_templates_support_enable_thinking(chat_params.tmpls.get()) : false;222 223        // Apply chat template to the list of messages224        return common_chat_templates_apply(chat_params.tmpls.get(), inputs);225    }226};227 228// TODO?: Make this reusable, enums, docs229static const std::array<const std::string, 7> cmds = {230    "/audio ",231    "/clear",232    "/exit",233    "/glob ",234    "/image ",235    "/read ",236    "/regen",237};238 239static std::vector<std::pair<std::string, size_t>> auto_completion_callback(std::string_view line, size_t cursor_byte_pos) {240    std::vector<std::pair<std::string, size_t>> matches;241    std::string cmd;242 243    if (line.length() > 1 && line[0] == '/' && !std::any_of(cmds.begin(), cmds.end(), [line](const std::string & prefix) {244        return string_starts_with(line, prefix);245    })) {246        auto it = cmds.begin();247 248        while ((it = std::find_if(it, cmds.end(), [line](const std::string & cmd_line) {249            return string_starts_with(cmd_line, line);250        })) != cmds.end()) {251            matches.emplace_back(*it, (*it).length());252            ++it;253        }254    } else {255        auto it = std::find_if(cmds.begin(), cmds.end(), [line](const std::string & prefix) {256            return prefix.back() == ' ' && string_starts_with(line, prefix);257        });258 259        if (it != cmds.end()) {260            cmd = *it;261        }262    }263 264    if (!cmd.empty() && cmd != "/glob " && line.length() >= cmd.length() && cursor_byte_pos >= cmd.length()) {265        const std::string path_prefix  = std::string(line.substr(cmd.length(), cursor_byte_pos - cmd.length()));266        const std::string path_postfix = std::string(line.substr(cursor_byte_pos));267        auto cur_dir = std::filesystem::current_path();268        std::string cur_dir_str = cur_dir.string();269        std::string expanded_prefix = path_prefix;270 271#if !defined(_WIN32)272        if (string_starts_with(path_prefix, "~")) {273            const char * home = std::getenv("HOME");274            if (home && home[0]) {275                expanded_prefix = std::string(home) + path_prefix.substr(1);276            }277        }278        if (string_starts_with(expanded_prefix, "/")) {279#else280        if (std::isalpha(expanded_prefix[0]) && expanded_prefix.find(':') == 1) {281#endif282            cur_dir = std::filesystem::path(expanded_prefix).parent_path();283            cur_dir_str = "";284        } else if (!path_prefix.empty()) {285            cur_dir /= std::filesystem::path(path_prefix).parent_path();286        }287 288        std::error_code ec;289        for (const auto & entry : std::filesystem::directory_iterator(cur_dir, ec)) {290            if (ec) {291                break;292            }293            if (!entry.exists(ec)) {294                ec.clear();295                continue;296            }297 298            const std::string path_full = entry.path().string();299            std::string path_entry = !cur_dir_str.empty() && string_starts_with(path_full, cur_dir_str) ? path_full.substr(cur_dir_str.length() + 1) : path_full;300 301            if (entry.is_directory(ec)) {302                path_entry.push_back(std::filesystem::path::preferred_separator);303            }304 305            if (expanded_prefix.empty() || string_starts_with(path_entry, expanded_prefix)) {306                std::string updated_line = cmd + path_entry;307                matches.emplace_back(updated_line + path_postfix, updated_line.length());308            }309 310            if (ec) {311                ec.clear();312            }313        }314 315        if (matches.empty()) {316            std::string updated_line = cmd + path_prefix;317            matches.emplace_back(updated_line + path_postfix, updated_line.length());318        }319 320        // Add the longest common prefix321        if (!expanded_prefix.empty() && matches.size() > 1) {322            const std::string_view match0(matches[0].first);323            const std::string_view match1(matches[1].first);324            auto it = std::mismatch(match0.begin(), match0.end(), match1.begin(), match1.end());325            size_t len = it.first - match0.begin();326 327            for (size_t i = 2; i < matches.size(); ++i) {328                const std::string_view matchi(matches[i].first);329                auto cmp = std::mismatch(match0.begin(), match0.end(), matchi.begin(), matchi.end());330                len = std::min(len, static_cast<size_t>(cmp.first - match0.begin()));331            }332 333            std::string updated_line = std::string(match0.substr(0, len));334            matches.emplace_back(updated_line + path_postfix, updated_line.length());335        }336 337        std::sort(matches.begin(), matches.end(), [](const auto & a, const auto & b) {338            return a.first.compare(0, a.second, b.first, 0, b.second) < 0;339        });340    }341 342    return matches;343}344 345static constexpr size_t FILE_GLOB_MAX_RESULTS = 100;346 347int main(int argc, char ** argv) {348    common_params params;349 350    params.verbosity = LOG_LEVEL_ERROR; // by default, less verbose logs351 352    common_init();353 354    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_CLI)) {355        return 1;356    }357 358    // TODO: maybe support it later?359    if (params.conversation_mode == COMMON_CONVERSATION_MODE_DISABLED) {360        console::error("--no-conversation is not supported by llama-cli\n");361        console::error("please use llama-completion instead\n");362    }363 364    // struct that contains llama context and inference365    cli_context ctx_cli(params);366 367    llama_backend_init();368    llama_numa_init(params.numa);369 370    // TODO: avoid using atexit() here by making `console` a singleton371    console::init(params.simple_io, params.use_color);372    atexit([]() { console::cleanup(); });373 374    console::set_display(DISPLAY_TYPE_RESET);375    console::set_completion_callback(auto_completion_callback);376 377#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))378    struct sigaction sigint_action;379    sigint_action.sa_handler = signal_handler;380    sigemptyset (&sigint_action.sa_mask);381    sigint_action.sa_flags = 0;382    sigaction(SIGINT, &sigint_action, NULL);383    sigaction(SIGTERM, &sigint_action, NULL);384#elif defined (_WIN32)385    auto console_ctrl_handler = +[](DWORD ctrl_type) -> BOOL {386        return (ctrl_type == CTRL_C_EVENT) ? (signal_handler(SIGINT), true) : false;387    };388    SetConsoleCtrlHandler(reinterpret_cast<PHANDLER_ROUTINE>(console_ctrl_handler), true);389#endif390 391    console::log("\nLoading model... "); // followed by loading animation392    console::spinner::start();393    if (!ctx_cli.ctx_server.load_model(params)) {394        console::spinner::stop();395        console::error("\nFailed to load the model\n");396        return 1;397    }398 399    console::spinner::stop();400    console::log("\n");401 402    std::thread inference_thread([&ctx_cli]() {403        ctx_cli.ctx_server.start_loop();404    });405 406    auto inf = ctx_cli.ctx_server.get_meta();407    std::string modalities = "text";408    if (inf.has_inp_image) {409        modalities += ", vision";410    }411    if (inf.has_inp_audio) {412        modalities += ", audio";413    }414 415    auto add_system_prompt = [&]() {416        if (!params.system_prompt.empty()) {417            ctx_cli.messages.push_back({418                {"role",    "system"},419                {"content", params.system_prompt}420            });421        }422    };423    add_system_prompt();424 425    console::log("\n");426    console::log("%s\n", LLAMA_ASCII_LOGO);427    console::log("build      : %s\n", inf.build_info.c_str());428    console::log("model      : %s\n", inf.model_name.c_str());429    console::log("modalities : %s\n", modalities.c_str());430    if (!params.system_prompt.empty()) {431        console::log("using custom system prompt\n");432    }433    console::log("\n");434    console::log("available commands:\n");435    console::log("  /exit or Ctrl+C     stop or exit\n");436    console::log("  /regen              regenerate the last response\n");437    console::log("  /clear              clear the chat history\n");438    console::log("  /read <file>        add a text file\n");439    console::log("  /glob <pattern>     add text files using globbing pattern\n");440    if (inf.has_inp_image) {441        console::log("  /image <file>       add an image file\n");442    }443    if (inf.has_inp_audio) {444        console::log("  /audio <file>       add an audio file\n");445    }446    console::log("\n");447 448    // interactive loop449    std::string cur_msg;450 451    auto add_text_file = [&](const std::string & fname) -> bool {452        std::string marker = ctx_cli.load_input_file(fname, false);453        if (marker.empty()) {454            console::error("file does not exist or cannot be opened: '%s'\n", fname.c_str());455            return false;456        }457        if (inf.fim_sep_token != LLAMA_TOKEN_NULL) {458            cur_msg += common_token_to_piece(ctx_cli.ctx_server.get_llama_context(), inf.fim_sep_token, true);459            cur_msg += fname;460            cur_msg.push_back('\n');461        } else {462            cur_msg += "--- File: ";463            cur_msg += fname;464            cur_msg += " ---\n";465        }466        cur_msg += marker;467        console::log("Loaded text from '%s'\n", fname.c_str());468        return true;469    };470 471    while (true) {472        std::string buffer;473        console::set_display(DISPLAY_TYPE_USER_INPUT);474        if (params.prompt.empty()) {475            console::log("\n> ");476            std::string line;477            bool another_line = true;478            do {479                another_line = console::readline(line, params.multiline_input);480                buffer += line;481            } while (another_line);482        } else {483            // process input prompt from args484            for (auto & fname : params.image) {485                std::string marker = ctx_cli.load_input_file(fname, true);486                if (marker.empty()) {487                    console::error("file does not exist or cannot be opened: '%s'\n", fname.c_str());488                    break;489                }490                console::log("Loaded media from '%s'\n", fname.c_str());491                cur_msg += marker;492            }493            buffer = params.prompt;494            if (buffer.size() > 500) {495                console::log("\n> %s ... (truncated)\n", buffer.substr(0, 500).c_str());496            } else {497                console::log("\n> %s\n", buffer.c_str());498            }499            params.prompt.clear(); // only use it once500        }501        console::set_display(DISPLAY_TYPE_RESET);502        console::log("\n");503 504        if (should_stop()) {505            g_is_interrupted.store(false);506            break;507        }508 509        // remove trailing newline510        if (!buffer.empty() &&buffer.back() == '\n') {511            buffer.pop_back();512        }513 514        // skip empty messages515        if (buffer.empty()) {516            continue;517        }518 519        bool add_user_msg = true;520 521        // process commands522        if (string_starts_with(buffer, "/exit")) {523            break;524        } else if (string_starts_with(buffer, "/regen")) {525            if (ctx_cli.messages.size() >= 2) {526                size_t last_idx = ctx_cli.messages.size() - 1;527                ctx_cli.messages.erase(last_idx);528                add_user_msg = false;529            } else {530                console::error("No message to regenerate.\n");531                continue;532            }533        } else if (string_starts_with(buffer, "/clear")) {534            ctx_cli.messages.clear();535            add_system_prompt();536 537            ctx_cli.input_files.clear();538            console::log("Chat history cleared.\n");539            continue;540        } else if (541                (string_starts_with(buffer, "/image ") && inf.has_inp_image) ||542                (string_starts_with(buffer, "/audio ") && inf.has_inp_audio)) {543            // just in case (bad copy-paste for example), we strip all trailing/leading spaces544            std::string fname = string_strip(buffer.substr(7));545            std::string marker = ctx_cli.load_input_file(fname, true);546            if (marker.empty()) {547                console::error("file does not exist or cannot be opened: '%s'\n", fname.c_str());548                continue;549            }550            cur_msg += marker;551            console::log("Loaded media from '%s'\n", fname.c_str());552            continue;553        } else if (string_starts_with(buffer, "/read ")) {554            std::string fname = string_strip(buffer.substr(6));555            add_text_file(fname);556            continue;557        } else if (string_starts_with(buffer, "/glob ")) {558            std::error_code ec;559            size_t count = 0;560            auto curdir = std::filesystem::current_path();561            std::string pattern = string_strip(buffer.substr(6));562            std::filesystem::path rel_path;563 564            auto startglob = pattern.find_first_of("![*?");565            if (startglob != std::string::npos && startglob != 0) {566                auto endpath = pattern.substr(0, startglob).find_last_of('/');567                if (endpath != std::string::npos) {568                    std::string rel_pattern = pattern.substr(0, endpath);569#if !defined(_WIN32)570                    if (string_starts_with(rel_pattern, "~")) {571                        const char * home = std::getenv("HOME");572                        if (home && home[0]) {573                            rel_pattern = std::string(home) + rel_pattern.substr(1);574                        }575                    }576#endif577                    rel_path = rel_pattern;578                    pattern.erase(0, endpath + 1);579                    curdir /= rel_path;580                }581            }582 583            for (const auto & entry : std::filesystem::recursive_directory_iterator(curdir,584                    std::filesystem::directory_options::skip_permission_denied, ec)) {585                if (!entry.is_regular_file()) {586                    continue;587                }588 589                std::string rel = std::filesystem::relative(entry.path(), curdir, ec).string();590                if (ec) {591                    ec.clear();592                    continue;593                }594                std::replace(rel.begin(), rel.end(), '\\', '/');595 596                if (!glob_match(pattern, rel)) {597                    continue;598                }599 600                if (!add_text_file((rel_path / rel).string())) {601                    continue;602                }603 604                if (++count >= FILE_GLOB_MAX_RESULTS) {605                    console::error("Maximum number of globbed files allowed (%zu) reached.\n", FILE_GLOB_MAX_RESULTS);606                    break;607                }608            }609            continue;610        } else {611            // not a command612            cur_msg += buffer;613        }614 615        // generate response616        if (add_user_msg) {617            ctx_cli.messages.push_back({618                {"role",    "user"},619                {"content", cur_msg}620            });621            cur_msg.clear();622        }623        result_timings timings;624        std::string assistant_content = ctx_cli.generate_completion(timings);625        ctx_cli.messages.push_back({626            {"role",    "assistant"},627            {"content", assistant_content}628        });629        console::log("\n");630 631        if (params.show_timings) {632            console::set_display(DISPLAY_TYPE_INFO);633            console::log("\n");634            console::log("[ Prompt: %.1f t/s | Generation: %.1f t/s ]\n", timings.prompt_per_second, timings.predicted_per_second);635            console::set_display(DISPLAY_TYPE_RESET);636        }637 638        if (params.single_turn) {639            break;640        }641    }642 643    console::set_display(DISPLAY_TYPE_RESET);644 645    console::log("\nExiting...\n");646    ctx_cli.ctx_server.terminate();647    inference_thread.join();648 649    // bump the log level to display timings650    common_log_set_verbosity_thold(LOG_LEVEL_INFO);651    common_memory_breakdown_print(ctx_cli.ctx_server.get_llama_context());652 653    return 0;654}655