llama-cli: add support for reasoning

common/console.cpp

@@ -1,6 +1,7 @@
 #include "console.h"
 #include <vector>
 #include <iostream>
+#include <cstdarg>
 
 #if defined(_WIN32)
 #define WIN32_LEAN_AND_MEAN
@@ -24,14 +25,14 @@
 #include <termios.h>
 #endif
 
-#define ANSI_COLOR_RED     "\x1b[31m"
-#define ANSI_COLOR_GREEN   "\x1b[32m"
-#define ANSI_COLOR_YELLOW  "\x1b[33m"
-#define ANSI_COLOR_BLUE    "\x1b[34m"
-#define ANSI_COLOR_MAGENTA "\x1b[35m"
-#define ANSI_COLOR_CYAN    "\x1b[36m"
-#define ANSI_COLOR_RESET   "\x1b[0m"
-#define ANSI_BOLD          "\x1b[1m"
+#define ANSI_COLOR_RED     LOG_COL_RED
+#define ANSI_COLOR_GREEN   LOG_COL_GREEN
+#define ANSI_COLOR_YELLOW  LOG_COL_YELLOW
+#define ANSI_COLOR_BLUE    LOG_COL_BLUE
+#define ANSI_COLOR_MAGENTA LOG_COL_MAGENTA
+#define ANSI_COLOR_CYAN    LOG_COL_CYAN
+#define ANSI_COLOR_RESET   LOG_COL_DEFAULT
+#define ANSI_BOLD          LOG_COL_BOLD
 
 namespace console {
 
@@ -142,25 +143,65 @@ namespace console {
     // Keep track of current display and only emit ANSI code if it changes
     void set_display(display_t display) {
         if (advanced_display && current_display != display) {
-            fflush(stdout);
-            switch(display) {
-                case reset:
-                    fprintf(out, ANSI_COLOR_RESET);
-                    break;
-                case prompt:
-                    fprintf(out, ANSI_COLOR_YELLOW);
-                    break;
-                case user_input:
-                    fprintf(out, ANSI_BOLD ANSI_COLOR_GREEN);
-                    break;
-                case error:
-                    fprintf(out, ANSI_BOLD ANSI_COLOR_RED);
-            }
             current_display = display;
-            fflush(out);
+
+            if (display == user_input && common_log_is_active(common_log_main())) {
+                common_log_flush(common_log_main());
+            }
+
+            if (display == user_input || !common_log_is_active(common_log_main())) {
+                fflush(stdout);
+                switch(display) {
+                    case reset:
+                        fprintf(out, ANSI_COLOR_RESET);
+                        break;
+                    case prompt:
+                        fprintf(out, ANSI_COLOR_YELLOW);
+                        break;
+                    case user_input:
+                        fprintf(out, ANSI_BOLD ANSI_COLOR_GREEN);
+                        break;
+                    case error:
+                        fprintf(out, ANSI_BOLD ANSI_COLOR_RED);
+                        break;
+                    case reasoning:
+                        fprintf(out, ANSI_COLOR_BLUE);
+                        break;
+                }
+                fflush(out);
+            }
         }
     }
 
+    display_t get_display() {
+        return current_display;
+    }
+
+    const char * get_display_color() {
+        switch(current_display) {
+            case reset:
+                return ANSI_COLOR_RESET;
+            case prompt:
+                return ANSI_COLOR_YELLOW;
+            case user_input:
+                return ANSI_BOLD ANSI_COLOR_GREEN;
+            case error:
+                return ANSI_BOLD ANSI_COLOR_RED;
+            case reasoning:
+                return ANSI_COLOR_BLUE;
+            default:
+                return "";
+        }
+    }
+
+    void write_console(const char * format, ...) {
+        va_list args;
+        va_start(args, format);
+        vfprintf(out, format, args);
+        va_end(args);
+        fflush(out);
+    }
+
     static char32_t getchar32() {
 #if defined(_WIN32)
         HANDLE hConsole = GetStdHandle(STD_INPUT_HANDLE);

common/console.h

@@ -3,17 +3,43 @@
 #pragma once
 
 #include <string>
+#include "log.h"
 
 namespace console {
     enum display_t {
         reset = 0,
         prompt,
         user_input,
-        error
+        error,
+        reasoning
     };
 
     void init(bool use_simple_io, bool use_advanced_display);
     void cleanup();
     void set_display(display_t display);
+    display_t    get_display();
+    const char * get_display_color();
     bool readline(std::string & line, bool multiline_input);
+
+    void write_console(const char * format, ...);
+
+    template<typename... Args>
+    void write(const char * format, Args... args) {
+        if (get_display() == user_input || !common_log_is_active(common_log_main())) {
+            write_console(format, args...);
+
+        } else {
+            const char * color = get_display_color();
+            std::string colored_format = std::string(color) + format + LOG_COL_DEFAULT;
+            common_log_add(common_log_main(), GGML_LOG_LEVEL_CONT, colored_format.c_str(), args...);
+        }
+    }
+
+    inline void write(const char * data) {
+        write("%s", data);
+    }
+
+    inline void write(const std::string & data) {
+        write("%s", data.c_str());
+    }
 }

common/log.cpp

@@ -174,6 +174,7 @@ struct common_log {
     std::mutex mtx;
     std::thread thrd;
     std::condition_variable cv;
+    std::condition_variable cv_flushed;
 
     FILE * file;
 
@@ -288,6 +289,10 @@ struct common_log {
                     cur = entries[head];
 
                     head = (head + 1) % entries.size();
+
+                    if (head == tail) {
+                        cv_flushed.notify_all();
+                    }
                 }
 
                 if (cur.is_end) {
@@ -376,6 +381,18 @@ struct common_log {
 
         this->timestamps = timestamps;
     }
+
+    bool is_active() const {
+        return running;
+    }
+
+    void flush() {
+        if (!running) {
+            return;
+        }
+        std::unique_lock<std::mutex> lock(mtx);
+        cv_flushed.wait(lock, [this]() { return head == tail; });
+    }
 };
 
 //
@@ -409,6 +426,14 @@ void common_log_free(struct common_log * log) {
     delete log;
 }
 
+bool common_log_is_active(struct common_log * log) {
+    return log->is_active();
+}
+
+void common_log_flush(struct common_log * log) {
+    log->flush();
+}
+
 void common_log_add(struct common_log * log, enum ggml_log_level level, const char * fmt, ...) {
     va_list args;
     va_start(args, fmt);

common/log.h

@@ -44,9 +44,11 @@ struct common_log;
 
 struct common_log * common_log_init();
 struct common_log * common_log_main(); // singleton, automatically destroys itself on exit
-void                common_log_pause (struct common_log * log); // pause  the worker thread, not thread-safe
-void                common_log_resume(struct common_log * log); // resume the worker thread, not thread-safe
-void                common_log_free  (struct common_log * log);
+void                common_log_pause    (struct common_log * log); // pause  the worker thread, not thread-safe
+void                common_log_resume   (struct common_log * log); // resume the worker thread, not thread-safe
+void                common_log_free     (struct common_log * log);
+bool                common_log_is_active(struct common_log * log); // check if logging is active
+void                common_log_flush    (struct common_log * log); // wait for all pending messages to be processed
 
 LOG_ATTRIBUTE_FORMAT(3, 4)
 void common_log_add(struct common_log * log, enum ggml_log_level level, const char * fmt, ...);