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)

---
 src/config/mod.rs | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

(limited to 'src/config')

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