1 files changed, 51 insertions, 94 deletions

diff --git a/src/client/azure_openai.rs b/src/client/azure_openai.rs
index 80d484d..45a5ecc 100644
--- a/src/client/azure_openai.rs
+++ b/src/client/azure_openai.rs
@@ -1,34 +1,33 @@
-use super::openai::{openai_send_message, openai_send_message_streaming};
-use super::{set_proxy, Client, ClientConfig, ModelInfo};
+use super::openai::{openai_build_body, openai_send_message, openai_send_message_streaming};
+use super::{
+    prompt_input_api_base, prompt_input_api_key, prompt_input_max_token, prompt_input_model_name,
+    Client, ClientConfig, ExtraConfig, ModelInfo, SendData,
+};
 
 use crate::config::SharedConfig;
 use crate::repl::ReplyStreamHandler;
 
-use anyhow::{anyhow, Context, Result};
+use anyhow::{anyhow, Result};
 use async_trait::async_trait;
-use inquire::{Confirm, Text};
 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)]
 #[derive(Debug)]
 pub struct AzureOpenAIClient {
     global_config: SharedConfig,
-    local_config: AzureOpenAIConfig,
+    config: AzureOpenAIConfig,
     model_info: ModelInfo,
 }
 
 #[derive(Debug, Clone, Deserialize)]
 pub struct AzureOpenAIConfig {
+    pub name: Option<String>,
     pub api_base: String,
     pub api_key: Option<String>,
     pub models: Vec<AzureOpenAIModel>,
-    pub proxy: Option<String>,
-    /// Set a timeout in seconds for connect to server
-    pub connect_timeout: Option<u64>,
+    pub extra: Option<ExtraConfig>,
 }
 
 #[derive(Debug, Clone, Deserialize)]
@@ -39,32 +38,36 @@ pub struct AzureOpenAIModel {
 
 #[async_trait]
 impl Client for AzureOpenAIClient {
-    fn get_config(&self) -> &SharedConfig {
+    fn config(&self) -> &SharedConfig {
         &self.global_config
     }
 
-    async fn send_message_inner(&self, content: &str) -> Result<String> {
-        let builder = self.request_builder(content, false)?;
+    fn extra_config(&self) -> &Option<ExtraConfig> {
+        &self.config.extra
+    }
+
+    async fn send_message_inner(&self, client: &ReqwestClient, data: SendData) -> Result<String> {
+        let builder = self.request_builder(client, data)?;
         openai_send_message(builder).await
     }
 
     async fn send_message_streaming_inner(
         &self,
-        content: &str,
+        client: &ReqwestClient,
         handler: &mut ReplyStreamHandler,
+        data: SendData,
     ) -> Result<()> {
-        let builder = self.request_builder(content, true)?;
+        let builder = self.request_builder(client, data)?;
         openai_send_message_streaming(builder, handler).await
     }
 }
 
 impl AzureOpenAIClient {
+    pub const NAME: &str = "azure-openai";
+
     pub fn init(global_config: SharedConfig) -> Option<Box<dyn Client>> {
         let model_info = global_config.read().model_info.clone();
-        if model_info.client != AzureOpenAIClient::name() {
-            return None;
-        }
-        let local_config = {
+        let config = {
             if let ClientConfig::AzureOpenAI(c) = &global_config.read().clients[model_info.index] {
                 c.clone()
             } else {
@@ -73,59 +76,37 @@ impl AzureOpenAIClient {
         };
         Some(Box::new(Self {
             global_config,
-            local_config,
+            config,
             model_info,
         }))
     }
 
-    pub fn name() -> &'static str {
-        "azure-openai"
+    pub fn name(local_config: &AzureOpenAIConfig) -> &str {
+        local_config.name.as_deref().unwrap_or(Self::NAME)
     }
 
     pub fn list_models(local_config: &AzureOpenAIConfig, index: usize) -> Vec<ModelInfo> {
+        let client = Self::name(local_config);
+
         local_config
             .models
             .iter()
-            .map(|v| ModelInfo::new(Self::name(), &v.name, v.max_tokens, index))
+            .map(|v| ModelInfo::new(client, &v.name, v.max_tokens, index))
             .collect()
     }
 
     pub fn create_config() -> Result<String> {
-        let mut client_config = format!("clients:\n  - type: {}\n", Self::name());
-
-        let api_base = Text::new("api_base:")
-            .prompt()
-            .map_err(|_| anyhow!("An error happened when asking for api base, try again later."))?;
+        let mut client_config = format!("clients:\n  - type: {}\n", Self::NAME);
 
+        let api_base = prompt_input_api_base()?;
         client_config.push_str(&format!("    api_base: {api_base}\n"));
 
-        if env::var("AZURE_OPENAI_KEY").is_err() {
-            let api_key = Text::new("API key:").prompt().map_err(|_| {
-                anyhow!("An error happened when asking for api key, try again later.")
-            })?;
+        let api_key = prompt_input_api_key()?;
+        client_config.push_str(&format!("    api_key: {api_key}\n"));
 
-            client_config.push_str(&format!("    api_key: {api_key}\n"));
-        }
+        let model_name = prompt_input_model_name()?;
 
-        let model_name = Text::new("Model Name:").prompt().map_err(|_| {
-            anyhow!("An error happened when asking for model name, try again later.")
-        })?;
-
-        let max_tokens = Text::new("Max tokens:").prompt().map_err(|_| {
-            anyhow!("An error happened when asking for max tokens, try again later.")
-        })?;
-
-        let ans = Confirm::new("Use proxy?")
-            .with_default(false)
-            .prompt()
-            .map_err(|_| anyhow!("Not finish questionnaire, try again later."))?;
-
-        if ans {
-            let proxy = Text::new("Set proxy:").prompt().map_err(|_| {
-                anyhow!("An error happened when asking for proxy, try again later.")
-            })?;
-            client_config.push_str(&format!("    proxy: {proxy}\n"));
-        }
+        let max_tokens = prompt_input_max_token()?;
 
         client_config.push_str(&format!(
             "    models:\n      - name: {model_name}\n        max_tokens: {max_tokens}\n"
@@ -134,50 +115,26 @@ impl AzureOpenAIClient {
         Ok(client_config)
     }
 
-    fn request_builder(&self, content: &str, stream: bool) -> Result<RequestBuilder> {
-        let messages = self.global_config.read().build_messages(content)?;
-
-        let mut body = json!({
-            "messages": messages,
-        });
-
-        if let Some(v) = self.global_config.read().get_temperature() {
-            body.as_object_mut()
-                .and_then(|m| m.insert("temperature".into(), json!(v)));
-        }
-
-        if stream {
-            body.as_object_mut()
-                .and_then(|m| m.insert("stream".into(), json!(true)));
-        }
-
-        let client = {
-            let mut builder = ReqwestClient::builder();
-            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)
-                .build()
-                .with_context(|| "Failed to build client")?
-        };
-        let mut api_base = self.local_config.api_base.clone();
-        if !api_base.ends_with('/') {
-            api_base = format!("{api_base}/");
-        }
+    fn request_builder(&self, client: &ReqwestClient, data: SendData) -> Result<RequestBuilder> {
+        let api_key = self.config.api_key.clone();
+        let api_key = api_key
+            .or_else(|| {
+                let env_prefix = match &self.config.name {
+                    None => "AZURE".into(),
+                    Some(v) => v.to_uppercase(),
+                };
+                env::var(format!("{env_prefix}_OPENAI_KEY")).ok()
+            })
+            .ok_or_else(|| anyhow!("Miss api_key"))?;
+
+        let body = openai_build_body(data, self.model_info.name.clone());
 
         let url = format!(
-            "{api_base}openai/deployments/{}/chat/completions?api-version=2023-05-15",
-            self.model_info.name
+            "{}/openai/deployments/{}/chat/completions?api-version=2023-05-15",
+            self.config.api_base, self.model_info.name
         );
 
-        let mut builder = client.post(url);
-
-        if let Some(api_key) = &self.local_config.api_key {
-            builder = builder.header("api-key", api_key)
-        } else if let Ok(api_key) = env::var("AZURE_OPENAI_KEY") {
-            builder = builder.header("api-key", api_key)
-        }
-        builder = builder.json(&body);
+        let builder = client.post(url).header("api-key", api_key).json(&body);
 
         Ok(builder)
     }