refactor: improve environment variables (#163)

Rename `AICHAT_API_KEY` to `OPENAI_API_KEY`
Add `LOCALAI_API_KEY`

3 files changed, 30 insertions, 16 deletions

diff --git a/src/client/localai.rs b/src/client/localai.rs
index 52d7ab3..5fe5f01 100644
--- a/src/client/localai.rs
+++ b/src/client/localai.rs
@@ -1,5 +1,5 @@
 use super::openai::{openai_send_message, openai_send_message_streaming};
-use super::{Client, ModelInfo};
+use super::{set_proxy, Client, ModelInfo};
 
 use crate::config::SharedConfig;
 use crate::repl::ReplyStreamHandler;
@@ -7,9 +7,10 @@ use crate::repl::ReplyStreamHandler;
 use anyhow::{anyhow, Context, Result};
 use async_trait::async_trait;
 use inquire::{Confirm, Text};
-use reqwest::{Client as ReqwestClient, Proxy, RequestBuilder};
+use reqwest::{Client as ReqwestClient, RequestBuilder};
 use serde::Deserialize;
 use serde_json::json;
+use std::env;
 use std::time::Duration;
 
 #[allow(clippy::module_name_repetitions)]
@@ -151,10 +152,7 @@ impl LocalAIClient {
 
         let client = {
             let mut builder = ReqwestClient::builder();
-            if let Some(proxy) = &self.local_config.proxy {
-                builder = builder
-                    .proxy(Proxy::all(proxy).with_context(|| format!("Invalid proxy `{proxy}`"))?);
-            }
+            builder = set_proxy(builder, &self.local_config.proxy)?;
             let timeout = Duration::from_secs(self.local_config.connect_timeout.unwrap_or(10));
             builder
                 .connect_timeout(timeout)
@@ -165,7 +163,9 @@ impl LocalAIClient {
         let mut builder = client.post(&self.local_config.url);
         if let Some(api_key) = &self.local_config.api_key {
             builder = builder.bearer_auth(api_key);
-        };
+        } else if let Ok(api_key) = env::var("LOCALAI_API_KEY") {
+            builder = builder.bearer_auth(api_key);
+        }
         builder = builder.json(&body);
 
         Ok(builder)
diff --git a/src/client/mod.rs b/src/client/mod.rs
index 1f9e66c..b9e478e 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -8,8 +8,9 @@ use self::{
 
 use anyhow::{bail, Context, Result};
 use async_trait::async_trait;
+use reqwest::{ClientBuilder, Proxy};
 use serde::Deserialize;
-use std::time::Duration;
+use std::{env, time::Duration};
 use tokio::runtime::Runtime;
 use tokio::time::sleep;
 
@@ -204,3 +205,19 @@ pub fn init_runtime() -> Result<Runtime> {
         .build()
         .with_context(|| "Failed to init tokio")
 }
+
+pub(crate) fn set_proxy(builder: ClientBuilder, proxy: &Option<String>) -> Result<ClientBuilder> {
+    let proxy = if let Some(proxy) = proxy {
+        if proxy.is_empty() || proxy == "false" || proxy == "-" {
+            return Ok(builder);
+        }
+        proxy.clone()
+    } else if let Ok(proxy) = env::var("HTTPS_PROXY").or_else(|_| env::var("ALL_PROXY")) {
+        proxy
+    } else {
+        return Ok(builder);
+    };
+    let builder =
+        builder.proxy(Proxy::all(&proxy).with_context(|| format!("Invalid proxy `{proxy}`"))?);
+    Ok(builder)
+}
diff --git a/src/client/openai.rs b/src/client/openai.rs
index 8ad3bab..885a02d 100644
--- a/src/client/openai.rs
+++ b/src/client/openai.rs
@@ -1,14 +1,14 @@
-use super::{Client, ModelInfo};
+use super::{set_proxy, Client, ModelInfo};
 
+use crate::config::SharedConfig;
 use crate::repl::ReplyStreamHandler;
-use crate::{config::SharedConfig, utils::get_env_name};
 
 use anyhow::{anyhow, bail, Context, Result};
 use async_trait::async_trait;
 use eventsource_stream::Eventsource;
 use futures_util::StreamExt;
 use inquire::{Confirm, Text};
-use reqwest::{Client as ReqwestClient, Proxy, RequestBuilder};
+use reqwest::{Client as ReqwestClient, RequestBuilder};
 use serde::Deserialize;
 use serde_json::{json, Value};
 use std::env;
@@ -120,7 +120,7 @@ impl OpenAIClient {
     fn request_builder(&self, content: &str, stream: bool) -> Result<RequestBuilder> {
         let api_key = if let Some(api_key) = &self.local_config.api_key {
             api_key.to_string()
-        } else if let Ok(api_key) = env::var(get_env_name("api_key")) {
+        } else if let Ok(api_key) = env::var("OPENAI_API_KEY") {
             api_key.to_string()
         } else {
             bail!("Miss api_key")
@@ -145,10 +145,7 @@ impl OpenAIClient {
 
         let client = {
             let mut builder = ReqwestClient::builder();
-            if let Some(proxy) = &self.local_config.proxy {
-                builder = builder
-                    .proxy(Proxy::all(proxy).with_context(|| format!("Invalid proxy `{proxy}`"))?);
-            }
+            builder = set_proxy(builder, &self.local_config.proxy)?;
             let timeout = Duration::from_secs(self.local_config.connect_timeout.unwrap_or(10));
             builder
                 .connect_timeout(timeout)