diff options
| author | sigoden <sigoden@gmail.com> | 2024-11-06 06:46:29 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-11-06 06:46:29 +0800 |
| commit | 35ee277a234a92add2f6244ad6f7ff3616b17ebe (patch) | |
| tree | b52de0378cb6427717027e7a4a5d74710088fb28 | |
| parent | 640eb00d74d42d98ab117e48bf4f181350bae221 (diff) | |
| download | aichat-35ee277a234a92add2f6244ad6f7ff3616b17ebe.tar.gz | |
feat: add config `user_agent` (#971)
| -rw-r--r-- | config.example.yaml | 1 | ||||
| -rw-r--r-- | src/client/common.rs | 3 | ||||
| -rw-r--r-- | src/config/mod.rs | 16 |
3 files changed, 20 insertions, 0 deletions
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<()> { |