feat: enhance logger (#731)

4 files changed, 73 insertions, 51 deletions

diff --git a/src/config/mod.rs b/src/config/mod.rs
index 8116532..7b22496 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -23,6 +23,7 @@ use inquire::{Confirm, Select};
 use parking_lot::RwLock;
 use serde::Deserialize;
 use serde_json::json;
+use simplelog::LevelFilter;
 use std::collections::{HashMap, HashSet};
 use std::{
     env,
@@ -386,6 +387,36 @@ impl Config {
         flags
     }
 
+    pub fn log(is_serve: bool) -> Result<(LevelFilter, Option<PathBuf>)> {
+        let log_level = env::var(get_env_name("log_level"))
+            .ok()
+            .and_then(|v| v.parse().ok())
+            .unwrap_or(match cfg!(debug_assertions) {
+                true => LevelFilter::Debug,
+                false => {
+                    if is_serve {
+                        LevelFilter::Info
+                    } else {
+                        LevelFilter::Off
+                    }
+                }
+            });
+        if log_level == LevelFilter::Off {
+            return Ok((log_level, None));
+        }
+        let log_path = match env::var(get_env_name("log_path")) {
+            Ok(v) => Some(PathBuf::from(v)),
+            Err(_) => match is_serve {
+                true => None,
+                false => Some(Config::local_path(&format!(
+                    "{}.log",
+                    env!("CARGO_CRATE_NAME")
+                ))?),
+            },
+        };
+        Ok((log_level, log_path))
+    }
+
     pub fn current_model(&self) -> &Model {
         if let Some(session) = self.session.as_ref() {
             session.model()
@@ -466,7 +497,7 @@ impl Config {
             .clone()
             .map_or_else(|| String::from("no"), |v| v.to_string());
         let role = self.extract_role();
-        let items = vec![
+        let mut items = vec![
             ("model", role.model().id()),
             (
                 "max_output_tokens",
@@ -509,6 +540,9 @@ impl Config {
             ("sessions_dir", display_path(&self.sessions_dir()?)),
             ("messages_file", display_path(&self.messages_file()?)),
         ];
+        if let Ok((_, Some(log_path))) = Self::log(self.working_mode.is_serve()) {
+            items.push(("log_path", display_path(&log_path)));
+        }
         let output = items
             .iter()
             .map(|(name, value)| format!("{name:<24}{value}"))
diff --git a/src/logger.rs b/src/logger.rs
deleted file mode 100644
index 25d6382..0000000
--- a/src/logger.rs
+++ /dev/null
@@ -1,45 +0,0 @@
-use crate::config::WorkingMode;
-
-use anyhow::Result;
-use log::LevelFilter;
-use simplelog::{format_description, Config as LogConfig, ConfigBuilder};
-
-#[cfg(debug_assertions)]
-pub fn setup_logger(working_mode: WorkingMode) -> Result<()> {
-    let is_serve = working_mode.is_serve();
-    let config = build_config(is_serve);
-    if is_serve {
-        simplelog::SimpleLogger::init(LevelFilter::Debug, config)?;
-    } else {
-        let file = std::fs::File::create(crate::config::Config::local_path("debug.log")?)?;
-        simplelog::WriteLogger::init(LevelFilter::Debug, config, file)?;
-    }
-    Ok(())
-}
-
-#[cfg(not(debug_assertions))]
-pub fn setup_logger(working_mode: WorkingMode) -> Result<()> {
-    if working_mode.is_serve() {
-        let config = build_config(true);
-        simplelog::SimpleLogger::init(log::LevelFilter::Info, config)?;
-    }
-    Ok(())
-}
-
-fn build_config(is_serve: bool) -> LogConfig {
-    let log_filter = if is_serve {
-        "aichat::serve".into()
-    } else {
-        match std::env::var("AICHAT_LOG_FILTER") {
-            Ok(v) => v,
-            Err(_) => "aichat".into(),
-        }
-    };
-    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()
-}
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(())
+}
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index e151745..5c48b1b 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -413,9 +413,10 @@ impl Repl {
     }
 
     fn banner(&self) {
+        let name = env!("CARGO_CRATE_NAME");
         let version = env!("CARGO_PKG_VERSION");
         print!(
-            r#"Welcome to aichat {version}
+            r#"Welcome to {name} {version}
 Type ".help" for additional help.
 "#
         )