From 341f2d994295d1b88bb849d74893d62df610150c Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Sun, 21 Jul 2024 17:24:04 +0800
Subject: feat: enhance logger (#731)

---
 src/main.rs | 40 ++++++++++++++++++++++++++++++++++++----
 1 file changed, 36 insertions(+), 4 deletions(-)

(limited to 'src/main.rs')

diff --git a/src/main.rs b/src/main.rs
index 13276a5..e2d86f5 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,7 +2,6 @@ mod cli;
 mod client;
 mod config;
 mod function;
-mod logger;
 mod rag;
 mod render;
 mod repl;
@@ -23,8 +22,8 @@ use crate::function::{eval_tool_calls, need_send_tool_results};
 use crate::render::render_error;
 use crate::repl::Repl;
 use crate::utils::{
-    create_abort_signal, create_spinner, detect_shell, extract_block, run_command, AbortSignal,
-    Shell, CODE_BLOCK_RE, IS_STDOUT_TERMINAL,
+    create_abort_signal, create_spinner, detect_shell, extract_block, get_env_name, run_command,
+    AbortSignal, Shell, CODE_BLOCK_RE, IS_STDOUT_TERMINAL,
 };
 
 use anyhow::{bail, Result};
@@ -33,6 +32,7 @@ use clap::Parser;
 use inquire::{Select, Text};
 use is_terminal::IsTerminal;
 use parking_lot::RwLock;
+use simplelog::{format_description, ConfigBuilder, LevelFilter, SimpleLogger, WriteLogger};
 use std::io::{stderr, stdin, Read};
 use std::process;
 use std::sync::Arc;
@@ -52,7 +52,7 @@ async fn main() -> Result<()> {
     } else {
         WorkingMode::Command
     };
-    crate::logger::setup_logger(working_mode)?;
+    setup_logger(working_mode.is_serve())?;
     let config = Arc::new(RwLock::new(Config::init(working_mode)?));
 
     let abort_signal = create_abort_signal();
@@ -313,3 +313,35 @@ async fn create_input(
     }
     Ok(input)
 }
+
+fn setup_logger(is_serve: bool) -> Result<()> {
+    let (log_level, log_path) = Config::log(is_serve)?;
+    if log_level == LevelFilter::Off {
+        return Ok(());
+    }
+    let crate_name = env!("CARGO_CRATE_NAME");
+    let log_filter = match std::env::var(get_env_name("log_filter")) {
+        Ok(v) => v,
+        Err(_) => match is_serve {
+            true => format!("{crate_name}::serve"),
+            false => crate_name.into(),
+        },
+    };
+    let config = ConfigBuilder::new()
+        .add_filter_allow(log_filter)
+        .set_time_format_custom(format_description!(
+            "[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z"
+        ))
+        .set_thread_level(LevelFilter::Off)
+        .build();
+    match log_path {
+        None => {
+            SimpleLogger::init(log_level, config)?;
+        }
+        Some(log_path) => {
+            let log_file = std::fs::File::create(log_path)?;
+            WriteLogger::init(log_level, config, log_file)?;
+        }
+    }
+    Ok(())
+}
-- 
cgit v1.2.3