refactor: clients/* and config.rs (#193)

- add register_clients macro to make it easier to add a new client
- no create_client_config, just add const PROMPTS
- move ModelInfo from clients/ to config/
- model's max_tokens are optional
- improve code quanity on config/mod.rs
- add/use macro config_get_fn

5 files changed, 405 insertions, 464 deletions

diff --git a/src/client/azure_openai.rs b/src/client/azure_openai.rs
index 45a5ecc..cabcebd 100644
--- a/src/client/azure_openai.rs
+++ b/src/client/azure_openai.rs
@@ -1,8 +1,5 @@
 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 super::{AzureOpenAIClient, Client, ExtraConfig, ModelInfo, PromptKind, PromptType, SendData};
 
 use crate::config::SharedConfig;
 use crate::repl::ReplyStreamHandler;
@@ -14,17 +11,10 @@ use serde::Deserialize;
 
 use std::env;
 
-#[derive(Debug)]
-pub struct AzureOpenAIClient {
-    global_config: SharedConfig,
-    config: AzureOpenAIConfig,
-    model_info: ModelInfo,
-}
-
 #[derive(Debug, Clone, Deserialize)]
 pub struct AzureOpenAIConfig {
     pub name: Option<String>,
-    pub api_base: String,
+    pub api_base: Option<String>,
     pub api_key: Option<String>,
     pub models: Vec<AzureOpenAIModel>,
     pub extra: Option<ExtraConfig>,
@@ -33,17 +23,13 @@ pub struct AzureOpenAIConfig {
 #[derive(Debug, Clone, Deserialize)]
 pub struct AzureOpenAIModel {
     name: String,
-    max_tokens: usize,
+    max_tokens: Option<usize>,
 }
 
 #[async_trait]
 impl Client for AzureOpenAIClient {
-    fn config(&self) -> &SharedConfig {
-        &self.global_config
-    }
-
-    fn extra_config(&self) -> &Option<ExtraConfig> {
-        &self.config.extra
+    fn config(&self) -> (&SharedConfig, &Option<ExtraConfig>) {
+        (&self.global_config, &self.config.extra)
     }
 
     async fn send_message_inner(&self, client: &ReqwestClient, data: SendData) -> Result<String> {
@@ -63,27 +49,19 @@ impl Client for AzureOpenAIClient {
 }
 
 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();
-        let config = {
-            if let ClientConfig::AzureOpenAI(c) = &global_config.read().clients[model_info.index] {
-                c.clone()
-            } else {
-                return None;
-            }
-        };
-        Some(Box::new(Self {
-            global_config,
-            config,
-            model_info,
-        }))
-    }
-
-    pub fn name(local_config: &AzureOpenAIConfig) -> &str {
-        local_config.name.as_deref().unwrap_or(Self::NAME)
-    }
+    config_get_fn!(api_base, get_api_base);
+
+    pub const PROMPTS: [PromptType<'static>; 4] = [
+        ("api_base", "API Base:", true, PromptKind::String),
+        ("api_key", "API Key:", true, PromptKind::String),
+        ("models[].name", "Model Name:", true, PromptKind::String),
+        (
+            "models[].max_tokens",
+            "Max Tokens:",
+            true,
+            PromptKind::Integer,
+        ),
+    ];
 
     pub fn list_models(local_config: &AzureOpenAIConfig, index: usize) -> Vec<ModelInfo> {
         let client = Self::name(local_config);
@@ -95,26 +73,6 @@ impl AzureOpenAIClient {
             .collect()
     }
 
-    pub fn create_config() -> Result<String> {
-        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"));
-
-        let api_key = prompt_input_api_key()?;
-        client_config.push_str(&format!("    api_key: {api_key}\n"));
-
-        let model_name = prompt_input_model_name()?;
-
-        let max_tokens = prompt_input_max_token()?;
-
-        client_config.push_str(&format!(
-            "    models:\n      - name: {model_name}\n        max_tokens: {max_tokens}\n"
-        ));
-
-        Ok(client_config)
-    }
-
     fn request_builder(&self, client: &ReqwestClient, data: SendData) -> Result<RequestBuilder> {
         let api_key = self.config.api_key.clone();
         let api_key = api_key
@@ -127,11 +85,13 @@ impl AzureOpenAIClient {
             })
             .ok_or_else(|| anyhow!("Miss api_key"))?;
 
+        let api_base = self.get_api_base()?;
+
         let body = openai_build_body(data, self.model_info.name.clone());
 
         let url = format!(
             "{}/openai/deployments/{}/chat/completions?api-version=2023-05-15",
-            self.config.api_base, self.model_info.name
+            &api_base, self.model_info.name
         );
 
         let builder = client.post(url).header("api-key", api_key).json(&body);
diff --git a/src/client/common.rs b/src/client/common.rs
new file mode 100644
index 0000000..6e5f365
--- /dev/null
+++ b/src/client/common.rs
@@ -0,0 +1,314 @@
+use crate::{
+    config::{Message, SharedConfig},
+    repl::{ReplyStreamHandler, SharedAbortSignal},
+    utils::{init_tokio_runtime, prompt_input_integer, prompt_input_string, tokenize, PromptKind},
+};
+
+use anyhow::{Context, Result};
+use async_trait::async_trait;
+use reqwest::{Client as ReqwestClient, ClientBuilder, Proxy};
+use serde::Deserialize;
+use serde_json::{json, Value};
+use std::{env, time::Duration};
+use tokio::time::sleep;
+
+use super::{openai::OpenAIConfig, ClientConfig};
+
+#[macro_export]
+macro_rules! register_role {
+    (
+        $(($name:literal, $config_key:ident, $config:ident, $client:ident),)+
+    ) => {
+
+        #[derive(Debug, Clone, Deserialize)]
+        #[serde(tag = "type")]
+        pub enum ClientConfig {
+            $(
+                #[serde(rename = $name)]
+                $config_key($config),
+            )+
+            #[serde(other)]
+            Unknown,
+        }
+
+
+        $(
+            #[derive(Debug)]
+            pub struct $client {
+                global_config: SharedConfig,
+                config: $config,
+                model_info: ModelInfo,
+            }
+
+            impl $client {
+                pub const NAME: &str = $name;
+
+                pub fn init(global_config: SharedConfig) -> Option<Box<dyn Client>> {
+                    let model_info = global_config.read().model_info.clone();
+                    let config = {
+                        if let ClientConfig::$config_key(c) = &global_config.read().clients[model_info.index] {
+                            c.clone()
+                        } else {
+                            return None;
+                        }
+                    };
+                    Some(Box::new(Self {
+                        global_config,
+                        config,
+                        model_info,
+                    }))
+                }
+
+                pub fn name(local_config: &$config) -> &str {
+                    local_config.name.as_deref().unwrap_or(Self::NAME)
+                }
+            }
+
+        )+
+
+        pub fn init_client(config: SharedConfig) -> Result<Box<dyn Client>> {
+                None
+                $(.or_else(|| $client::init(config.clone())))+
+                .ok_or_else(|| {
+                    let model_info = config.read().model_info.clone();
+                    anyhow!(
+                        "Unknown client {} at config.clients[{}]",
+                        &model_info.client,
+                        &model_info.index
+                    )
+                })
+        }
+
+        pub fn list_client_types() -> Vec<&'static str> {
+            vec![$($client::NAME,)+]
+        }
+
+        pub fn create_client_config(client: &str) -> Result<Value> {
+            $(
+                if client == $client::NAME {
+                    return create_config(&$client::PROMPTS, $client::NAME)
+                }
+            )+
+            bail!("Unknown client {}", client)
+        }
+
+        pub fn all_models(config: &Config) -> Vec<ModelInfo> {
+            config
+                .clients
+                .iter()
+                .enumerate()
+                .flat_map(|(i, v)| match v {
+                    $(ClientConfig::$config_key(c) => $client::list_models(c, i),)+
+                    ClientConfig::Unknown => vec![],
+                })
+                .collect()
+        }
+
+    };
+}
+
+macro_rules! config_get_fn {
+    ($field_name:ident, $fn_name:ident) => {
+        fn $fn_name(&self) -> Result<String> {
+            let api_key = self.config.$field_name.clone();
+            api_key
+                .or_else(|| {
+                    let env_prefix = Self::name(&self.config);
+                    let env_name =
+                        format!("{}_{}", env_prefix, stringify!($field_name)).to_ascii_uppercase();
+                    env::var(&env_name).ok()
+                })
+                .ok_or_else(|| anyhow::anyhow!("Miss {}", stringify!($field_name)))
+        }
+    };
+}
+
+#[async_trait]
+pub trait Client {
+    fn config(&self) -> (&SharedConfig, &Option<ExtraConfig>);
+
+    fn build_client(&self) -> Result<ReqwestClient> {
+        let mut builder = ReqwestClient::builder();
+        let options = self.config().1;
+        let timeout = options
+            .as_ref()
+            .and_then(|v| v.connect_timeout)
+            .unwrap_or(10);
+        let proxy = options.as_ref().and_then(|v| v.proxy.clone());
+        builder = set_proxy(builder, &proxy)?;
+        let client = builder
+            .connect_timeout(Duration::from_secs(timeout))
+            .build()
+            .with_context(|| "Failed to build client")?;
+        Ok(client)
+    }
+
+    fn send_message(&self, content: &str) -> Result<String> {
+        init_tokio_runtime()?.block_on(async {
+            let global_config = self.config().0;
+            if global_config.read().dry_run {
+                let content = global_config.read().echo_messages(content);
+                return Ok(content);
+            }
+            let client = self.build_client()?;
+            let data = global_config.read().prepare_send_data(content, false)?;
+            self.send_message_inner(&client, data)
+                .await
+                .with_context(|| "Failed to get awswer")
+        })
+    }
+
+    fn send_message_streaming(
+        &self,
+        content: &str,
+        handler: &mut ReplyStreamHandler,
+    ) -> Result<()> {
+        async fn watch_abort(abort: SharedAbortSignal) {
+            loop {
+                if abort.aborted() {
+                    break;
+                }
+                sleep(Duration::from_millis(100)).await;
+            }
+        }
+        let abort = handler.get_abort();
+        init_tokio_runtime()?.block_on(async {
+            tokio::select! {
+                ret = async {
+                    let global_config = self.config().0;
+                    if global_config.read().dry_run {
+                        let content = global_config.read().echo_messages(content);
+                        let tokens = tokenize(&content);
+                        for token in tokens {
+                            tokio::time::sleep(Duration::from_millis(25)).await;
+                            handler.text(&token)?;
+                        }
+                        return Ok(());
+                    }
+                    let client = self.build_client()?;
+                    let data = global_config.read().prepare_send_data(content, true)?;
+                    self.send_message_streaming_inner(&client, handler, data).await
+                } => {
+                    handler.done()?;
+                    ret.with_context(|| "Failed to get awswer")
+                }
+                _ = watch_abort(abort.clone()) => {
+                    handler.done()?;
+                    Ok(())
+                 },
+                _ =  tokio::signal::ctrl_c() => {
+                    abort.set_ctrlc();
+                    Ok(())
+                }
+            }
+        })
+    }
+
+    async fn send_message_inner(&self, client: &ReqwestClient, data: SendData) -> Result<String>;
+
+    async fn send_message_streaming_inner(
+        &self,
+        client: &ReqwestClient,
+        handler: &mut ReplyStreamHandler,
+        data: SendData,
+    ) -> Result<()>;
+}
+
+impl Default for ClientConfig {
+    fn default() -> Self {
+        Self::OpenAI(OpenAIConfig::default())
+    }
+}
+
+#[derive(Debug, Clone, Deserialize, Default)]
+pub struct ExtraConfig {
+    pub proxy: Option<String>,
+    pub connect_timeout: Option<u64>,
+}
+
+#[derive(Debug)]
+pub struct SendData {
+    pub messages: Vec<Message>,
+    pub temperature: Option<f64>,
+    pub stream: bool,
+}
+
+pub type PromptType<'a> = (&'a str, &'a str, bool, PromptKind);
+
+pub fn create_config(list: &[PromptType], client: &str) -> Result<Value> {
+    let mut config = json!({
+        "type": client,
+    });
+    for (path, desc, required, kind) in list {
+        match kind {
+            PromptKind::String => {
+                let value = prompt_input_string(desc, *required)?;
+                set_config_value(&mut config, path, kind, &value);
+            }
+            PromptKind::Integer => {
+                let value = prompt_input_integer(desc, *required)?;
+                set_config_value(&mut config, path, kind, &value);
+            }
+        }
+    }
+
+    let clients = json!(vec![config]);
+    Ok(clients)
+}
+
+fn set_config_value(json: &mut Value, path: &str, kind: &PromptKind, value: &str) {
+    let segs: Vec<&str> = path.split('.').collect();
+    match segs.as_slice() {
+        [name] => json[name] = to_json(kind, value),
+        [scope, name] => match scope.split_once('[') {
+            None => {
+                if json.get(scope).is_none() {
+                    let mut obj = json!({});
+                    obj[name] = to_json(kind, value);
+                    json[scope] = obj;
+                } else {
+                    json[scope][name] = to_json(kind, value);
+                }
+            }
+            Some((scope, _)) => {
+                if json.get(scope).is_none() {
+                    let mut obj = json!({});
+                    obj[name] = to_json(kind, value);
+                    json[scope] = json!([obj]);
+                } else {
+                    json[scope][0][name] = to_json(kind, value);
+                }
+            }
+        },
+        _ => {}
+    }
+}
+
+fn to_json(kind: &PromptKind, value: &str) -> Value {
+    if value.is_empty() {
+        return Value::Null;
+    }
+    match kind {
+        PromptKind::String => value.into(),
+        PromptKind::Integer => match value.parse::<i32>() {
+            Ok(value) => value.into(),
+            Err(_) => value.into(),
+        },
+    }
+}
+
+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/localai.rs b/src/client/localai.rs
index fb5a353..131d92b 100644
--- a/src/client/localai.rs
+++ b/src/client/localai.rs
@@ -1,8 +1,5 @@
 use super::openai::{openai_build_body, openai_send_message, openai_send_message_streaming};
-use super::{
-    prompt_input_api_base, prompt_input_api_key_optional, prompt_input_max_token,
-    prompt_input_model_name, Client, ClientConfig, ExtraConfig, ModelInfo, SendData,
-};
+use super::{Client, ExtraConfig, LocalAIClient, ModelInfo, PromptKind, PromptType, SendData};
 
 use crate::config::SharedConfig;
 use crate::repl::ReplyStreamHandler;
@@ -13,13 +10,6 @@ use reqwest::{Client as ReqwestClient, RequestBuilder};
 use serde::Deserialize;
 use std::env;
 
-#[derive(Debug)]
-pub struct LocalAIClient {
-    global_config: SharedConfig,
-    config: LocalAIConfig,
-    model_info: ModelInfo,
-}
-
 #[derive(Debug, Clone, Deserialize)]
 pub struct LocalAIConfig {
     pub name: Option<String>,
@@ -33,17 +23,13 @@ pub struct LocalAIConfig {
 #[derive(Debug, Clone, Deserialize)]
 pub struct LocalAIModel {
     name: String,
-    max_tokens: usize,
+    max_tokens: Option<usize>,
 }
 
 #[async_trait]
 impl Client for LocalAIClient {
-    fn config(&self) -> &SharedConfig {
-        &self.global_config
-    }
-
-    fn extra_config(&self) -> &Option<ExtraConfig> {
-        &self.config.extra
+    fn config(&self) -> (&SharedConfig, &Option<ExtraConfig>) {
+        (&self.global_config, &self.config.extra)
     }
 
     async fn send_message_inner(&self, client: &ReqwestClient, data: SendData) -> Result<String> {
@@ -63,27 +49,19 @@ impl Client for LocalAIClient {
 }
 
 impl LocalAIClient {
-    pub const NAME: &str = "localai";
-
-    pub fn init(global_config: SharedConfig) -> Option<Box<dyn Client>> {
-        let model_info = global_config.read().model_info.clone();
-        let config = {
-            if let ClientConfig::LocalAI(c) = &global_config.read().clients[model_info.index] {
-                c.clone()
-            } else {
-                return None;
-            }
-        };
-        Some(Box::new(Self {
-            global_config,
-            config,
-            model_info,
-        }))
-    }
-
-    pub fn name(local_config: &LocalAIConfig) -> &str {
-        local_config.name.as_deref().unwrap_or(Self::NAME)
-    }
+    config_get_fn!(api_key, get_api_key);
+
+    pub const PROMPTS: [PromptType<'static>; 4] = [
+        ("api_base", "API Base:", true, PromptKind::String),
+        ("api_key", "API Key:", false, PromptKind::String),
+        ("models[].name", "Model Name:", true, PromptKind::String),
+        (
+            "models[].max_tokens",
+            "Max Tokens:",
+            false,
+            PromptKind::Integer,
+        ),
+    ];
 
     pub fn list_models(local_config: &LocalAIConfig, index: usize) -> Vec<ModelInfo> {
         let client = Self::name(local_config);
@@ -95,32 +73,8 @@ impl LocalAIClient {
             .collect()
     }
 
-    pub fn create_config() -> Result<String> {
-        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"));
-
-        let api_key = prompt_input_api_key_optional()?;
-        client_config.push_str(&format!("    api_key: {api_key}\n"));
-
-        let model_name = prompt_input_model_name()?;
-
-        let max_tokens = prompt_input_max_token()?;
-
-        client_config.push_str(&format!(
-            "    models:\n      - name: {model_name}\n        max_tokens: {max_tokens}\n"
-        ));
-
-        Ok(client_config)
-    }
-
     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 = Self::name(&self.config).to_uppercase();
-            env::var(format!("{env_prefix}_API_KEY")).ok()
-        });
+        let api_key = self.get_api_key().ok();
 
         let body = openai_build_body(data, self.model_info.name.clone());
 
diff --git a/src/client/mod.rs b/src/client/mod.rs
index bef6094..5562eb6 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -1,272 +1,32 @@
+#[macro_use]
+mod common;
+
 pub mod azure_openai;
 pub mod localai;
 pub mod openai;
 
-use self::{
-    azure_openai::{AzureOpenAIClient, AzureOpenAIConfig},
-    localai::LocalAIConfig,
-    openai::{OpenAIClient, OpenAIConfig},
-};
+pub use common::*;
+
+use self::azure_openai::AzureOpenAIConfig;
+use self::localai::LocalAIConfig;
+use self::openai::OpenAIConfig;
 
 use crate::{
-    client::localai::LocalAIClient,
-    config::{Config, Message, SharedConfig},
-    repl::{ReplyStreamHandler, SharedAbortSignal},
-    utils::tokenize,
+    config::{Config, ModelInfo, SharedConfig},
+    utils::PromptKind,
 };
 
-use anyhow::{anyhow, bail, Context, Result};
-use async_trait::async_trait;
-use inquire::{required, Text};
-use reqwest::{Client as ReqwestClient, ClientBuilder, Proxy};
+use anyhow::{anyhow, bail, Result};
 use serde::Deserialize;
-use std::{env, time::Duration};
-use tokio::time::sleep;
-
-#[derive(Debug, Clone, Deserialize)]
-#[serde(tag = "type")]
-pub enum ClientConfig {
-    #[serde(rename = "openai")]
-    OpenAI(OpenAIConfig),
-    #[serde(rename = "localai")]
-    LocalAI(LocalAIConfig),
-    #[serde(rename = "azure-openai")]
-    AzureOpenAI(AzureOpenAIConfig),
-}
-#[derive(Debug, Clone)]
-pub struct ModelInfo {
-    pub client: String,
-    pub name: String,
-    pub max_tokens: usize,
-    pub index: usize,
-}
-
-impl Default for ModelInfo {
-    fn default() -> Self {
-        OpenAIClient::list_models(&OpenAIConfig::default(), 0)[0].clone()
-    }
-}
-
-impl ModelInfo {
-    pub fn new(client: &str, name: &str, max_tokens: usize, index: usize) -> Self {
-        Self {
-            client: client.into(),
-            name: name.into(),
-            max_tokens,
-            index,
-        }
-    }
-    pub fn stringify(&self) -> String {
-        format!("{}:{}", self.client, self.name)
-    }
-}
-
-#[derive(Debug)]
-pub struct SendData {
-    pub messages: Vec<Message>,
-    pub temperature: Option<f64>,
-    pub stream: bool,
-}
-#[async_trait]
-pub trait Client {
-    fn config(&self) -> &SharedConfig;
-
-    fn extra_config(&self) -> &Option<ExtraConfig>;
-
-    fn build_client(&self) -> Result<ReqwestClient> {
-        let mut builder = ReqwestClient::builder();
-        let options = self.extra_config();
-        let timeout = options
-            .as_ref()
-            .and_then(|v| v.connect_timeout)
-            .unwrap_or(10);
-        let proxy = options.as_ref().and_then(|v| v.proxy.clone());
-        builder = set_proxy(builder, &proxy)?;
-        let client = builder
-            .connect_timeout(Duration::from_secs(timeout))
-            .build()
-            .with_context(|| "Failed to build client")?;
-        Ok(client)
-    }
-
-    fn send_message(&self, content: &str) -> Result<String> {
-        init_tokio_runtime()?.block_on(async {
-            if self.config().read().dry_run {
-                let content = self.config().read().echo_messages(content);
-                return Ok(content);
-            }
-            let client = self.build_client()?;
-            let data = self.config().read().prepare_send_data(content, false)?;
-            self.send_message_inner(&client, data)
-                .await
-                .with_context(|| "Failed to fetch")
-        })
-    }
-
-    fn send_message_streaming(
-        &self,
-        content: &str,
-        handler: &mut ReplyStreamHandler,
-    ) -> Result<()> {
-        async fn watch_abort(abort: SharedAbortSignal) {
-            loop {
-                if abort.aborted() {
-                    break;
-                }
-                sleep(Duration::from_millis(100)).await;
-            }
-        }
-        let abort = handler.get_abort();
-        init_tokio_runtime()?.block_on(async {
-            tokio::select! {
-                ret = async {
-                    if self.config().read().dry_run {
-                        let content = self.config().read().echo_messages(content);
-                        let tokens = tokenize(&content);
-                        for token in tokens {
-                            tokio::time::sleep(Duration::from_millis(25)).await;
-                            handler.text(&token)?;
-                        }
-                        return Ok(());
-                    }
-                    let client = self.build_client()?;
-                    let data = self.config().read().prepare_send_data(content, true)?;
-                    self.send_message_streaming_inner(&client, handler, data).await
-                } => {
-                    handler.done()?;
-                    ret.with_context(|| "Failed to fetch stream")
-                }
-                _ = watch_abort(abort.clone()) => {
-                    handler.done()?;
-                    Ok(())
-                 },
-                _ =  tokio::signal::ctrl_c() => {
-                    abort.set_ctrlc();
-                    Ok(())
-                }
-            }
-        })
-    }
-
-    async fn send_message_inner(&self, client: &ReqwestClient, data: SendData) -> Result<String>;
-
-    async fn send_message_streaming_inner(
-        &self,
-        client: &ReqwestClient,
-        handler: &mut ReplyStreamHandler,
-        data: SendData,
-    ) -> Result<()>;
-}
-
-#[derive(Debug, Clone, Deserialize, Default)]
-pub struct ExtraConfig {
-    pub proxy: Option<String>,
-    pub connect_timeout: Option<u64>,
-}
-
-pub fn init_client(config: SharedConfig) -> Result<Box<dyn Client>> {
-    OpenAIClient::init(config.clone())
-        .or_else(|| LocalAIClient::init(config.clone()))
-        .or_else(|| AzureOpenAIClient::init(config.clone()))
-        .ok_or_else(|| {
-            let model_info = config.read().model_info.clone();
-            anyhow!(
-                "Unknown client {} at config.clients[{}]",
-                &model_info.client,
-                &model_info.index
-            )
-        })
-}
-
-pub fn list_client_types() -> Vec<&'static str> {
-    vec![
-        OpenAIClient::NAME,
-        LocalAIClient::NAME,
-        AzureOpenAIClient::NAME,
-    ]
-}
-
-pub fn create_client_config(client: &str) -> Result<String> {
-    if client == OpenAIClient::NAME {
-        OpenAIClient::create_config()
-    } else if client == LocalAIClient::NAME {
-        LocalAIClient::create_config()
-    } else if client == AzureOpenAIClient::NAME {
-        AzureOpenAIClient::create_config()
-    } else {
-        bail!("Unknown client {}", &client)
-    }
-}
-
-pub fn list_models(config: &Config) -> Vec<ModelInfo> {
-    config
-        .clients
-        .iter()
-        .enumerate()
-        .flat_map(|(i, v)| match v {
-            ClientConfig::OpenAI(c) => OpenAIClient::list_models(c, i),
-            ClientConfig::LocalAI(c) => LocalAIClient::list_models(c, i),
-            ClientConfig::AzureOpenAI(c) => AzureOpenAIClient::list_models(c, i),
-        })
-        .collect()
-}
-
-pub(crate) fn init_tokio_runtime() -> Result<tokio::runtime::Runtime> {
-    tokio::runtime::Builder::new_current_thread()
-        .enable_all()
-        .build()
-        .with_context(|| "Failed to init tokio")
-}
-
-pub(crate) fn prompt_input_api_base() -> Result<String> {
-    Text::new("API Base:")
-        .with_validator(required!("This field is required"))
-        .prompt()
-        .map_err(prompt_op_err)
-}
-
-pub(crate) fn prompt_input_api_key() -> Result<String> {
-    Text::new("API Key:")
-        .with_validator(required!("This field is required"))
-        .prompt()
-        .map_err(prompt_op_err)
-}
-
-pub(crate) fn prompt_input_api_key_optional() -> Result<String> {
-    Text::new("API Key:").prompt().map_err(prompt_op_err)
-}
-
-pub(crate) fn prompt_input_model_name() -> Result<String> {
-    Text::new("Model Name:")
-        .with_validator(required!("This field is required"))
-        .prompt()
-        .map_err(prompt_op_err)
-}
-
-pub(crate) fn prompt_input_max_token() -> Result<String> {
-    Text::new("Max tokens:")
-        .with_default("4096")
-        .with_validator(required!("This field is required"))
-        .prompt()
-        .map_err(prompt_op_err)
-}
-
-pub(crate) fn prompt_op_err<T>(_: T) -> anyhow::Error {
-    anyhow!("An error happened, try again later.")
-}
-
-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)
-}
+use serde_json::Value;
+
+register_role!(
+    ("openai", OpenAI, OpenAIConfig, OpenAIClient),
+    ("localai", LocalAI, LocalAIConfig, LocalAIClient),
+    (
+        "azure-openai",
+        AzureOpenAI,
+        AzureOpenAIConfig,
+        AzureOpenAIClient
+    ),
+);
diff --git a/src/client/openai.rs b/src/client/openai.rs
index 57cd5fd..80969e4 100644
--- a/src/client/openai.rs
+++ b/src/client/openai.rs
@@ -1,4 +1,4 @@
-use super::{prompt_input_api_key, Client, ClientConfig, ExtraConfig, ModelInfo, SendData};
+use super::{Client, ExtraConfig, ModelInfo, OpenAIClient, PromptKind, PromptType, SendData};
 
 use crate::config::SharedConfig;
 use crate::repl::ReplyStreamHandler;
@@ -14,12 +14,12 @@ use std::env;
 
 const API_BASE: &str = "https://api.openai.com/v1";
 
-#[derive(Debug)]
-pub struct OpenAIClient {
-    global_config: SharedConfig,
-    config: OpenAIConfig,
-    model_info: ModelInfo,
-}
+const MODELS: [(&str, usize); 4] = [
+    ("gpt-3.5-turbo", 4096),
+    ("gpt-3.5-turbo-16k", 16384),
+    ("gpt-4", 8192),
+    ("gpt-4-32k", 32768),
+];
 
 #[derive(Debug, Clone, Deserialize, Default)]
 pub struct OpenAIConfig {
@@ -31,12 +31,8 @@ pub struct OpenAIConfig {
 
 #[async_trait]
 impl Client for OpenAIClient {
-    fn config(&self) -> &SharedConfig {
-        &self.global_config
-    }
-
-    fn extra_config(&self) -> &Option<ExtraConfig> {
-        &self.config.extra
+    fn config(&self) -> (&SharedConfig, &Option<ExtraConfig>) {
+        (&self.global_config, &self.config.extra)
     }
 
     async fn send_message_inner(&self, client: &ReqwestClient, data: SendData) -> Result<String> {
@@ -56,61 +52,25 @@ impl Client for OpenAIClient {
 }
 
 impl OpenAIClient {
-    pub const NAME: &str = "openai";
-
-    pub fn init(global_config: SharedConfig) -> Option<Box<dyn Client>> {
-        let model_info = global_config.read().model_info.clone();
-        let config = {
-            if let ClientConfig::OpenAI(c) = &global_config.read().clients[model_info.index] {
-                c.clone()
-            } else {
-                return None;
-            }
-        };
-        Some(Box::new(Self {
-            global_config,
-            config,
-            model_info,
-        }))
-    }
+    config_get_fn!(api_key, get_api_key);
 
-    pub fn name(local_config: &OpenAIConfig) -> &str {
-        local_config.name.as_deref().unwrap_or(Self::NAME)
-    }
+    pub const PROMPTS: [PromptType<'static>; 1] =
+        [("api_key", "API Key:", true, PromptKind::String)];
 
     pub fn list_models(local_config: &OpenAIConfig, index: usize) -> Vec<ModelInfo> {
         let client = Self::name(local_config);
-
-        [
-            ("gpt-3.5-turbo", 4096),
-            ("gpt-3.5-turbo-16k", 16384),
-            ("gpt-4", 8192),
-            ("gpt-4-32k", 32768),
-        ]
-        .into_iter()
-        .map(|(name, max_tokens)| ModelInfo::new(client, name, max_tokens, index))
-        .collect()
-    }
-
-    pub fn create_config() -> Result<String> {
-        let mut client_config = format!("clients:\n  - type: {}\n", Self::NAME);
-
-        let api_key = prompt_input_api_key()?;
-        client_config.push_str(&format!("    api_key: {api_key}\n"));
-
-        Ok(client_config)
+        MODELS
+            .into_iter()
+            .map(|(name, max_tokens)| ModelInfo::new(client, name, Some(max_tokens), index))
+            .collect()
     }
 
     fn request_builder(&self, client: &ReqwestClient, data: SendData) -> Result<RequestBuilder> {
-        let env_prefix = Self::name(&self.config).to_uppercase();
-
-        let api_key = self.config.api_key.clone();
-        let api_key = api_key
-            .or_else(|| env::var(format!("{env_prefix}_API_KEY")).ok())
-            .ok_or_else(|| anyhow!("Miss api_key"))?;
+        let api_key = self.get_api_key()?;
 
         let body = openai_build_body(data, self.model_info.name.clone());
 
+        let env_prefix = Self::name(&self.config).to_uppercase();
         let api_base = env::var(format!("{env_prefix}_API_BASE"))
             .ok()
             .unwrap_or_else(|| API_BASE.to_string());
@@ -127,20 +87,20 @@ impl OpenAIClient {
     }
 }
 
-pub(crate) async fn openai_send_message(builder: RequestBuilder) -> Result<String> {
+pub async fn openai_send_message(builder: RequestBuilder) -> Result<String> {
     let data: Value = builder.send().await?.json().await?;
     if let Some(err_msg) = data["error"]["message"].as_str() {
-        bail!("Request failed, {err_msg}");
+        bail!("{err_msg}");
     }
 
     let output = data["choices"][0]["message"]["content"]
         .as_str()
-        .ok_or_else(|| anyhow!("Unexpected response {data}"))?;
+        .ok_or_else(|| anyhow!("Invalid response data: {data}"))?;
 
     Ok(output.to_string())
 }
 
-pub(crate) async fn openai_send_message_streaming(
+pub async fn openai_send_message_streaming(
     builder: RequestBuilder,
     handler: &mut ReplyStreamHandler,
 ) -> Result<()> {
@@ -148,7 +108,7 @@ pub(crate) async fn openai_send_message_streaming(
     if !res.status().is_success() {
         let data: Value = res.json().await?;
         if let Some(err_msg) = data["error"]["message"].as_str() {
-            bail!("Request failed, {err_msg}");
+            bail!("{err_msg}");
         }
         bail!("Request failed");
     }
@@ -159,37 +119,30 @@ pub(crate) async fn openai_send_message_streaming(
             break;
         }
         let data: Value = serde_json::from_str(&chunk)?;
-        let text = data["choices"][0]["delta"]["content"]
-            .as_str()
-            .unwrap_or_default();
-        if text.is_empty() {
-            continue;
+        if let Some(text) = data["choices"][0]["delta"]["content"].as_str() {
+            handler.text(text)?;
         }
-        handler.text(text)?;
     }
 
     Ok(())
 }
 
-pub(crate) fn openai_build_body(data: SendData, model: String) -> Value {
+pub fn openai_build_body(data: SendData, model: String) -> Value {
     let SendData {
         messages,
         temperature,
         stream,
     } = data;
+
     let mut body = json!({
         "model": model,
         "messages": messages,
     });
-
     if let Some(v) = temperature {
-        body.as_object_mut()
-            .and_then(|m| m.insert("temperature".into(), json!(v)));
+        body["temperature"] = v.into();
     }
-
     if stream {
-        body.as_object_mut()
-            .and_then(|m| m.insert("stream".into(), json!(true)));
+        body["stream"] = true.into();
     }
     body
 }