From 35ee277a234a92add2f6244ad6f7ff3616b17ebe Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Wed, 6 Nov 2024 06:46:29 +0800
Subject: feat: add config `user_agent` (#971)

---
 config.example.yaml  |  1 +
 src/client/common.rs |  3 +++
 src/config/mod.rs    | 16 ++++++++++++++++
 3 files changed, 20 insertions(+)

diff --git a/config.example.yaml b/config.example.yaml
index 6ea3ea5..7e6cdd1 100644
--- a/config.example.yaml
+++ b/config.example.yaml
@@ -82,6 +82,7 @@ right_prompt:
 
 # ---- misc ----
 serve_addr: 127.0.0.1:8000                  # Default serve listening address 
+use_agent: null                             # Set Use-Agent HTTP header, use `auto` for aichat/<current-version>
 
 # ---- clients ----
 clients:
diff --git a/src/client/common.rs b/src/client/common.rs
index 1275d0b..9c516c8 100644
--- a/src/client/common.rs
+++ b/src/client/common.rs
@@ -43,6 +43,9 @@ pub trait Client: Sync + Send {
         let timeout = extra.and_then(|v| v.connect_timeout).unwrap_or(10);
         let proxy = extra.and_then(|v| v.proxy.clone());
         builder = set_proxy(builder, proxy.as_ref())?;
+        if let Some(user_agent) = self.global_config().read().user_agent.as_ref() {
+            builder = builder.user_agent(user_agent);
+        }
         let client = builder
             .connect_timeout(Duration::from_secs(timeout))
             .build()
diff --git a/src/config/mod.rs b/src/config/mod.rs
index adf36f3..124eaa5 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -134,6 +134,7 @@ pub struct Config {
     pub right_prompt: Option<String>,
 
     pub serve_addr: Option<String>,
+    pub user_agent: Option<String>,
 
     pub clients: Vec<ClientConfig>,
 
@@ -200,6 +201,7 @@ impl Default for Config {
             right_prompt: None,
 
             serve_addr: None,
+            user_agent: None,
 
             clients: vec![],
 
@@ -246,6 +248,7 @@ impl Config {
 
         config.setup_model()?;
         config.setup_document_loaders();
+        config.setup_user_agent();
 
         Ok(config)
     }
@@ -2068,6 +2071,9 @@ impl Config {
         if let Some(v) = read_env_value::<String>("serve_addr") {
             self.serve_addr = v;
         }
+        if let Some(v) = read_env_value::<String>("user_agent") {
+            self.user_agent = v;
+        }
     }
 
     fn load_functions(&mut self) -> Result<()> {
@@ -2097,6 +2103,16 @@ impl Config {
                 self.document_loaders.entry(k).or_insert(v);
             });
     }
+
+    fn setup_user_agent(&mut self) {
+        if let Some("auto") = self.user_agent.as_deref() {
+            self.user_agent = Some(format!(
+                "{}/{}",
+                env!("CARGO_CRATE_NAME"),
+                env!("CARGO_PKG_VERSION")
+            ));
+        }
+    }
 }
 
 pub fn load_env_file() -> Result<()> {
-- 
cgit v1.2.3