From a137483b03191d169e62f2049fcbdaace0821d6c Mon Sep 17 00:00:00 2001 From: sigoden Date: Tue, 31 Oct 2023 16:50:08 +0800 Subject: feat: improve client (#189) --- src/client/mod.rs | 126 ++++++++++++++++++++++++++++++++++++++++++------------ 1 file changed, 99 insertions(+), 27 deletions(-) (limited to 'src/client/mod.rs') diff --git a/src/client/mod.rs b/src/client/mod.rs index ca37804..bef6094 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -8,20 +8,21 @@ use self::{ openai::{OpenAIClient, OpenAIConfig}, }; -use anyhow::{anyhow, bail, Context, Result}; -use async_trait::async_trait; -use reqwest::{ClientBuilder, Proxy}; -use serde::Deserialize; -use std::{env, time::Duration}; -use tokio::time::sleep; - use crate::{ client::localai::LocalAIClient, - config::{Config, SharedConfig}, + config::{Config, Message, SharedConfig}, repl::{ReplyStreamHandler, SharedAbortSignal}, utils::tokenize, }; +use anyhow::{anyhow, bail, Context, Result}; +use async_trait::async_trait; +use inquire::{required, Text}; +use reqwest::{Client as ReqwestClient, ClientBuilder, Proxy}; +use serde::Deserialize; +use std::{env, time::Duration}; +use tokio::time::sleep; + #[derive(Debug, Clone, Deserialize)] #[serde(tag = "type")] pub enum ClientConfig { @@ -32,7 +33,6 @@ pub enum ClientConfig { #[serde(rename = "azure-openai")] AzureOpenAI(AzureOpenAIConfig), } - #[derive(Debug, Clone)] pub struct ModelInfo { pub client: String, @@ -61,17 +61,43 @@ impl ModelInfo { } } +#[derive(Debug)] +pub struct SendData { + pub messages: Vec, + pub temperature: Option, + pub stream: bool, +} #[async_trait] pub trait Client { - fn get_config(&self) -> &SharedConfig; + fn config(&self) -> &SharedConfig; + + fn extra_config(&self) -> &Option; + + fn build_client(&self) -> Result { + 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 { init_tokio_runtime()?.block_on(async { - if self.get_config().read().dry_run { - let content = self.get_config().read().echo_messages(content); + if self.config().read().dry_run { + let content = self.config().read().echo_messages(content); return Ok(content); } - self.send_message_inner(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") }) @@ -94,8 +120,8 @@ pub trait Client { init_tokio_runtime()?.block_on(async { tokio::select! { ret = async { - if self.get_config().read().dry_run { - let content = self.get_config().read().echo_messages(content); + 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; @@ -103,7 +129,9 @@ pub trait Client { } return Ok(()); } - self.send_message_streaming_inner(content, handler).await + 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") @@ -120,15 +148,22 @@ pub trait Client { }) } - async fn send_message_inner(&self, content: &str) -> Result; + async fn send_message_inner(&self, client: &ReqwestClient, data: SendData) -> Result; async fn send_message_streaming_inner( &self, - content: &str, + client: &ReqwestClient, handler: &mut ReplyStreamHandler, + data: SendData, ) -> Result<()>; } +#[derive(Debug, Clone, Deserialize, Default)] +pub struct ExtraConfig { + pub proxy: Option, + pub connect_timeout: Option, +} + pub fn init_client(config: SharedConfig) -> Result> { OpenAIClient::init(config.clone()) .or_else(|| LocalAIClient::init(config.clone())) @@ -143,20 +178,20 @@ pub fn init_client(config: SharedConfig) -> Result> { }) } -pub fn all_clients() -> Vec<&'static str> { +pub fn list_client_types() -> Vec<&'static str> { vec![ - OpenAIClient::name(), - LocalAIClient::name(), - AzureOpenAIClient::name(), + OpenAIClient::NAME, + LocalAIClient::NAME, + AzureOpenAIClient::NAME, ] } pub fn create_client_config(client: &str) -> Result { - if client == OpenAIClient::name() { + if client == OpenAIClient::NAME { OpenAIClient::create_config() - } else if client == LocalAIClient::name() { + } else if client == LocalAIClient::NAME { LocalAIClient::create_config() - } else if client == AzureOpenAIClient::name() { + } else if client == AzureOpenAIClient::NAME { AzureOpenAIClient::create_config() } else { bail!("Unknown client {}", &client) @@ -176,14 +211,51 @@ pub fn list_models(config: &Config) -> Vec { .collect() } -pub fn init_tokio_runtime() -> Result { +pub(crate) fn init_tokio_runtime() -> Result { tokio::runtime::Builder::new_current_thread() .enable_all() .build() .with_context(|| "Failed to init tokio") } -pub(crate) fn set_proxy(builder: ClientBuilder, proxy: &Option) -> Result { +pub(crate) fn prompt_input_api_base() -> Result { + 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 { + 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 { + Text::new("API Key:").prompt().map_err(prompt_op_err) +} + +pub(crate) fn prompt_input_model_name() -> Result { + 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 { + 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) -> anyhow::Error { + anyhow!("An error happened, try again later.") +} + +fn set_proxy(builder: ClientBuilder, proxy: &Option) -> Result { let proxy = if let Some(proxy) = proxy { if proxy.is_empty() || proxy == "false" || proxy == "-" { return Ok(builder); -- cgit v1.2.3