feat: support multi bots and custom url (#150)

3 files changed, 598 insertions, 0 deletions

diff --git a/src/client/localai.rs b/src/client/localai.rs
new file mode 100644
index 0000000..4375c1d
--- /dev/null
+++ b/src/client/localai.rs
@@ -0,0 +1,181 @@
+use super::openai::{openai_send_message, openai_send_message_streaming};
+use super::{Client, ModelInfo};
+
+use crate::config::SharedConfig;
+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 serde::Deserialize;
+use serde_json::json;
+use std::time::Duration;
+use tokio::runtime::Runtime;
+
+#[allow(clippy::module_name_repetitions)]
+#[derive(Debug)]
+pub struct LocalAIClient {
+    global_config: SharedConfig,
+    local_config: LocalAIConfig,
+    model_info: ModelInfo,
+    runtime: Runtime,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct LocalAIConfig {
+    pub url: String,
+    pub api_key: Option<String>,
+    pub models: Vec<LocalAIModel>,
+    pub proxy: Option<String>,
+    /// Set a timeout in seconds for connect to server
+    pub connect_timeout: Option<u64>,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct LocalAIModel {
+    name: String,
+    max_tokens: usize,
+}
+
+#[async_trait]
+impl Client for LocalAIClient {
+    fn get_config(&self) -> &SharedConfig {
+        &self.global_config
+    }
+
+    fn get_runtime(&self) -> &Runtime {
+        &self.runtime
+    }
+
+    async fn send_message_inner(&self, content: &str) -> Result<String> {
+        let builder = self.request_builder(content, false)?;
+        openai_send_message(builder).await
+    }
+
+    async fn send_message_streaming_inner(
+        &self,
+        content: &str,
+        handler: &mut ReplyStreamHandler,
+    ) -> Result<()> {
+        let builder = self.request_builder(content, true)?;
+        openai_send_message_streaming(builder, handler).await
+    }
+}
+
+impl LocalAIClient {
+    pub fn new(
+        global_config: SharedConfig,
+        local_config: LocalAIConfig,
+        model_info: ModelInfo,
+        runtime: Runtime,
+    ) -> Self {
+        Self {
+            global_config,
+            local_config,
+            model_info,
+            runtime,
+        }
+    }
+
+    pub fn name() -> &'static str {
+        "localai"
+    }
+
+    pub fn list_models(local_config: &LocalAIConfig) -> Vec<(String, usize)> {
+        local_config
+            .models
+            .iter()
+            .map(|v| (v.name.to_string(), v.max_tokens))
+            .collect()
+    }
+
+    pub fn create_config() -> Result<String> {
+        let mut client_config = format!("clients:\n  - type: {}\n", Self::name());
+
+        let url = Text::new("URL:")
+            .prompt()
+            .map_err(|_| anyhow!("An error happened when asking for url, try again later."))?;
+
+        client_config.push_str(&format!("    url: {url}\n"));
+
+        let ans = Confirm::new("Use auth?")
+            .with_default(false)
+            .prompt()
+            .map_err(|_| anyhow!("Not finish questionnaire, try again later."))?;
+
+        if ans {
+            let api_key = Text::new("API key:").prompt().map_err(|_| {
+                anyhow!("An error happened when asking for api key, try again later.")
+            })?;
+
+            client_config.push_str(&format!("    api_key: {api_key}\n"));
+        }
+
+        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"));
+        }
+
+        client_config.push_str(&format!(
+            "    models:\n      - name: {model_name}\n        max_tokens: {max_tokens}\n"
+        ));
+
+        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!({
+            "model": self.model_info.name,
+            "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();
+            if let Some(proxy) = &self.local_config.proxy {
+                builder = builder
+                    .proxy(Proxy::all(proxy).with_context(|| format!("Invalid proxy `{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 builder = client.post(&self.local_config.url);
+        if let Some(api_key) = &self.local_config.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
new file mode 100644
index 0000000..b541cff
--- /dev/null
+++ b/src/client/mod.rs
@@ -0,0 +1,198 @@
+pub mod localai;
+pub mod openai;
+
+use self::{
+    localai::LocalAIConfig,
+    openai::{OpenAIClient, OpenAIConfig},
+};
+
+use anyhow::{bail, Context, Result};
+use async_trait::async_trait;
+use serde::Deserialize;
+use std::time::Duration;
+use tokio::runtime::Runtime;
+use tokio::time::sleep;
+
+use crate::{
+    client::localai::LocalAIClient,
+    config::{Config, SharedConfig},
+    repl::{ReplyStreamHandler, SharedAbortSignal},
+};
+
+#[allow(clippy::struct_excessive_bools)]
+#[derive(Debug, Clone, Deserialize)]
+#[serde(tag = "type")]
+pub enum ClientConfig {
+    #[serde(rename = "openai")]
+    OpenAI(OpenAIConfig),
+    #[serde(rename = "localai")]
+    LocalAI(LocalAIConfig),
+}
+
+#[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 {
+        let client = OpenAIClient::name();
+        let (name, max_tokens) = &OpenAIClient::list_models(&OpenAIConfig::default())[0];
+        Self::new(client, name, *max_tokens, 0)
+    }
+}
+
+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)
+    }
+}
+
+#[async_trait]
+pub trait Client {
+    fn get_config(&self) -> &SharedConfig;
+
+    fn get_runtime(&self) -> &Runtime;
+
+    fn send_message(&self, content: &str) -> Result<String> {
+        self.get_runtime().block_on(async {
+            if self.get_config().read().dry_run {
+                return Ok(self.get_config().read().echo_messages(content));
+            }
+            self.send_message_inner(content)
+                .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();
+        self.get_runtime().block_on(async {
+            tokio::select! {
+                ret = async {
+                    if self.get_config().read().dry_run {
+                        handler.text(&self.get_config().read().echo_messages(content))?;
+                        return Ok(());
+                    }
+                    self.send_message_streaming_inner(content, handler).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, content: &str) -> Result<String>;
+
+    async fn send_message_streaming_inner(
+        &self,
+        content: &str,
+        handler: &mut ReplyStreamHandler,
+    ) -> Result<()>;
+}
+
+pub fn init_client(config: SharedConfig, runtime: Runtime) -> Result<Box<dyn Client>> {
+    let model_info = config.read().model_info.clone();
+    let model_info_err = |model_info: &ModelInfo| {
+        bail!(
+            "Unknown client {} at config.clients[{}]",
+            &model_info.client,
+            &model_info.index
+        )
+    };
+    if model_info.client == OpenAIClient::name() {
+        let local_config = {
+            if let ClientConfig::OpenAI(c) = &config.read().clients[model_info.index] {
+                c.clone()
+            } else {
+                return model_info_err(&model_info);
+            }
+        };
+        Ok(Box::new(OpenAIClient::new(
+            config,
+            local_config,
+            model_info,
+            runtime,
+        )))
+    } else if model_info.client == LocalAIClient::name() {
+        let local_config = {
+            if let ClientConfig::LocalAI(c) = &config.read().clients[model_info.index] {
+                c.clone()
+            } else {
+                return model_info_err(&model_info);
+            }
+        };
+        Ok(Box::new(LocalAIClient::new(
+            config,
+            local_config,
+            model_info,
+            runtime,
+        )))
+    } else {
+        bail!("Unknown client {}", &model_info.client)
+    }
+}
+
+pub fn all_clients() -> Vec<&'static str> {
+    vec![OpenAIClient::name(), LocalAIClient::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 {
+        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)
+                .iter()
+                .map(|(x, y)| ModelInfo::new(OpenAIClient::name(), x, *y, i))
+                .collect::<Vec<ModelInfo>>(),
+            ClientConfig::LocalAI(c) => LocalAIClient::list_models(c)
+                .iter()
+                .map(|(x, y)| ModelInfo::new(LocalAIClient::name(), x, *y, i))
+                .collect::<Vec<ModelInfo>>(),
+        })
+        .collect()
+}
diff --git a/src/client/openai.rs b/src/client/openai.rs
new file mode 100644
index 0000000..932dc79
--- /dev/null
+++ b/src/client/openai.rs
@@ -0,0 +1,219 @@
+use super::{Client, ModelInfo};
+
+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 serde::Deserialize;
+use serde_json::{json, Value};
+use std::env;
+use std::time::Duration;
+use tokio::runtime::Runtime;
+
+const API_URL: &str = "https://api.openai.com/v1/chat/completions";
+
+#[allow(clippy::module_name_repetitions)]
+#[derive(Debug)]
+pub struct OpenAIClient {
+    global_config: SharedConfig,
+    local_config: OpenAIConfig,
+    model_info: ModelInfo,
+    runtime: Runtime,
+}
+
+#[allow(clippy::struct_excessive_bools)]
+#[derive(Debug, Clone, Deserialize, Default)]
+pub struct OpenAIConfig {
+    pub api_key: Option<String>,
+    pub organization_id: Option<String>,
+    pub proxy: Option<String>,
+    /// Set a timeout in seconds for connect to openai server
+    pub connect_timeout: Option<u64>,
+}
+
+#[async_trait]
+impl Client for OpenAIClient {
+    fn get_config(&self) -> &SharedConfig {
+        &self.global_config
+    }
+
+    fn get_runtime(&self) -> &Runtime {
+        &self.runtime
+    }
+
+    async fn send_message_inner(&self, content: &str) -> Result<String> {
+        let builder = self.request_builder(content, false)?;
+        openai_send_message(builder).await
+    }
+
+    async fn send_message_streaming_inner(
+        &self,
+        content: &str,
+        handler: &mut ReplyStreamHandler,
+    ) -> Result<()> {
+        let builder = self.request_builder(content, true)?;
+        openai_send_message_streaming(builder, handler).await
+    }
+}
+
+impl OpenAIClient {
+    pub fn new(
+        global_config: SharedConfig,
+        local_config: OpenAIConfig,
+        model_info: ModelInfo,
+        runtime: Runtime,
+    ) -> Self {
+        Self {
+            global_config,
+            local_config,
+            model_info,
+            runtime,
+        }
+    }
+
+    pub fn name() -> &'static str {
+        "openai"
+    }
+
+    pub fn list_models(_local_config: &OpenAIConfig) -> Vec<(String, usize)> {
+        vec![
+            ("gpt-3.5-turbo".into(), 4096),
+            ("gpt-3.5-turbo-16k".into(), 16384),
+            ("gpt-4".into(), 8192),
+            ("gpt-4-32k".into(), 32768),
+        ]
+    }
+
+    pub fn create_config() -> Result<String> {
+        let mut client_config = format!("clients:\n  - type: {}\n", Self::name());
+
+        let api_key = Text::new("API key:")
+            .prompt()
+            .map_err(|_| anyhow!("An error happened when asking for api key, try again later."))?;
+
+        client_config.push_str(&format!("    api_key: {api_key}\n"));
+
+        let ans = Confirm::new("Has Organization?")
+            .with_default(false)
+            .prompt()
+            .map_err(|_| anyhow!("Not finish questionnaire, try again later."))?;
+
+        if ans {
+            let organization_id = Text::new("Organization ID:").prompt().map_err(|_| {
+                anyhow!("An error happened when asking for proxy, try again later.")
+            })?;
+            client_config.push_str(&format!("    organization_id: {organization_id}\n"));
+        }
+
+        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"));
+        }
+
+        Ok(client_config)
+    }
+
+    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")) {
+            api_key.to_string()
+        } else {
+            bail!("Miss api_key")
+        };
+
+        let messages = self.global_config.read().build_messages(content)?;
+
+        let mut body = json!({
+            "model": self.model_info.name,
+            "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();
+            if let Some(proxy) = &self.local_config.proxy {
+                builder = builder
+                    .proxy(Proxy::all(proxy).with_context(|| format!("Invalid proxy `{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 builder = client.post(API_URL).bearer_auth(api_key).json(&body);
+
+        if let Some(organization_id) = &self.local_config.organization_id {
+            builder = builder.header("OpenAI-Organization", organization_id);
+        }
+
+        Ok(builder)
+    }
+}
+
+pub(crate) 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}");
+    }
+
+    let output = data["choices"][0]["message"]["content"]
+        .as_str()
+        .ok_or_else(|| anyhow!("Unexpected response {data}"))?;
+
+    Ok(output.to_string())
+}
+
+pub(crate) async fn openai_send_message_streaming(
+    builder: RequestBuilder,
+    handler: &mut ReplyStreamHandler,
+) -> Result<()> {
+    let res = builder.send().await?;
+    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!("Request failed");
+    }
+    let mut stream = res.bytes_stream().eventsource();
+    while let Some(part) = stream.next().await {
+        let chunk = part?.data;
+        if chunk == "[DONE]" {
+            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;
+        }
+        handler.text(text)?;
+    }
+
+    Ok(())
+}