From 30afa19bc0e2be56be49476f4628a35314756fab Mon Sep 17 00:00:00 2001
From: Meng Zhang <meng@tabbyml.com>
Date: Tue, 12 Sep 2023 22:41:39 +0800
Subject: [PATCH] feat: add LLAMA_CPP_LOG_LEVEL to control log level of
 llama.cpp (#436)

---
 crates/llama-cpp-bindings/src/engine.cc | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/crates/llama-cpp-bindings/src/engine.cc b/crates/llama-cpp-bindings/src/engine.cc
index f250f62..7981bcf 100644
--- a/crates/llama-cpp-bindings/src/engine.cc
+++ b/crates/llama-cpp-bindings/src/engine.cc
@@ -100,8 +100,21 @@ class TextInferenceEngineImpl : public TextInferenceEngine {
   owned<llama_context> ctx_;
 };
 
+static int g_llama_cpp_log_level = 0;
+static void llama_log_callback(llama_log_level level, const char * text, void * user_data) {
+  (void)user_data;
+  if (level < g_llama_cpp_log_level) {
+    fputs(text, stderr);
+    fflush(stderr);
+  }
+}
+
 struct BackendInitializer {
   BackendInitializer() {
+    if (const char* level = std::getenv("LLAMA_CPP_LOG_LEVEL")) {
+      g_llama_cpp_log_level = std::stoi(level);
+    }
+    llama_log_set(llama_log_callback, nullptr);
     llama_backend_init(false);
   }