From 2ab2e23bb0539a7fac0d990e3d7e08b7fade1030 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Sat, 28 Oct 2023 22:10:17 +0800
Subject: refactor: improve environment variables (#163)

Rename `AICHAT_API_KEY` to `OPENAI_API_KEY`
Add `LOCALAI_API_KEY`
---
 src/client/mod.rs | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

(limited to 'src/client/mod.rs')

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)
+}
-- 
cgit v1.2.3