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/azure_openai.rs | 145 ++++++++++++++++---------------------------- src/client/localai.rs | 139 ++++++++++++++++-------------------------- src/client/mod.rs | 126 +++++++++++++++++++++++++++++--------- src/client/openai.rs | 146 +++++++++++++++++++-------------------------- 4 files changed, 261 insertions(+), 295 deletions(-) (limited to 'src/client') 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, pub api_base: String, pub api_key: Option, pub models: Vec, - pub proxy: Option, - /// Set a timeout in seconds for connect to server - pub connect_timeout: Option, + pub extra: Option, } #[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 { - let builder = self.request_builder(content, false)?; + fn extra_config(&self) -> &Option { + &self.config.extra + } + + async fn send_message_inner(&self, client: &ReqwestClient, data: SendData) -> Result { + 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> { 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 { + 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 { - 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 { - 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 { + 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) } diff --git a/src/client/localai.rs b/src/client/localai.rs index 084603b..fb5a353 100644 --- a/src/client/localai.rs +++ b/src/client/localai.rs @@ -1,33 +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_optional, 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::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; #[derive(Debug)] pub struct LocalAIClient { global_config: SharedConfig, - local_config: LocalAIConfig, + config: LocalAIConfig, model_info: ModelInfo, } #[derive(Debug, Clone, Deserialize)] pub struct LocalAIConfig { - pub url: String, + pub name: Option, + pub api_base: String, pub api_key: Option, + pub chat_endpoint: Option, pub models: Vec, - pub proxy: Option, - /// Set a timeout in seconds for connect to server - pub connect_timeout: Option, + pub extra: Option, } #[derive(Debug, Clone, Deserialize)] @@ -38,32 +38,36 @@ pub struct LocalAIModel { #[async_trait] impl Client for LocalAIClient { - fn get_config(&self) -> &SharedConfig { + fn config(&self) -> &SharedConfig { &self.global_config } - async fn send_message_inner(&self, content: &str) -> Result { - let builder = self.request_builder(content, false)?; + fn extra_config(&self) -> &Option { + &self.config.extra + } + + async fn send_message_inner(&self, client: &ReqwestClient, data: SendData) -> Result { + 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 LocalAIClient { + pub const NAME: &str = "localai"; + pub fn init(global_config: SharedConfig) -> Option> { let model_info = global_config.read().model_info.clone(); - if model_info.client != LocalAIClient::name() { - return None; - } - let local_config = { + let config = { if let ClientConfig::LocalAI(c) = &global_config.read().clients[model_info.index] { c.clone() } else { @@ -72,64 +76,37 @@ impl LocalAIClient { }; Some(Box::new(Self { global_config, - local_config, + config, model_info, })) } - pub fn name() -> &'static str { - "localai" + pub fn name(local_config: &LocalAIConfig) -> &str { + local_config.name.as_deref().unwrap_or(Self::NAME) } pub fn list_models(local_config: &LocalAIConfig, index: usize) -> Vec { + 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 { - let mut client_config = format!("clients:\n - type: {}\n", Self::name()); + 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."))?; + let api_base = prompt_input_api_base()?; + client_config.push_str(&format!(" api_base: {api_base}\n")); - 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 api_key = prompt_input_api_key_optional()?; + 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 model_name = prompt_input_model_name()?; - 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" @@ -138,41 +115,27 @@ impl LocalAIClient { Ok(client_config) } - fn request_builder(&self, content: &str, stream: bool) -> Result { - let messages = self.global_config.read().build_messages(content)?; - - let mut body = json!({ - "model": self.model_info.name, - "messages": messages, + fn request_builder(&self, client: &ReqwestClient, data: SendData) -> Result { + 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() }); - if let Some(v) = self.global_config.read().get_temperature() { - body.as_object_mut() - .and_then(|m| m.insert("temperature".into(), json!(v))); - } + let body = openai_build_body(data, self.model_info.name.clone()); - if stream { - body.as_object_mut() - .and_then(|m| m.insert("stream".into(), json!(true))); - } + let chat_endpoint = self + .config + .chat_endpoint + .as_deref() + .unwrap_or("/chat/completions"); - 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 url = format!("{}{chat_endpoint}", self.config.api_base); - 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); - } else if let Ok(api_key) = env::var("LOCALAI_API_KEY") { + let mut builder = client.post(url).json(&body); + if let Some(api_key) = 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 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); diff --git a/src/client/openai.rs b/src/client/openai.rs index e0d321d..57cd5fd 100644 --- a/src/client/openai.rs +++ b/src/client/openai.rs @@ -1,65 +1,66 @@ -use super::{set_proxy, Client, ClientConfig, ModelInfo}; +use super::{prompt_input_api_key, Client, ClientConfig, ExtraConfig, ModelInfo, SendData}; use crate::config::SharedConfig; use crate::repl::ReplyStreamHandler; -use anyhow::{anyhow, bail, Context, Result}; +use anyhow::{anyhow, bail, Result}; use async_trait::async_trait; use eventsource_stream::Eventsource; use futures_util::StreamExt; -use inquire::{Confirm, Text}; use reqwest::{Client as ReqwestClient, RequestBuilder}; use serde::Deserialize; use serde_json::{json, Value}; use std::env; -use std::time::Duration; const API_BASE: &str = "https://api.openai.com/v1"; #[derive(Debug)] pub struct OpenAIClient { global_config: SharedConfig, - local_config: OpenAIConfig, + config: OpenAIConfig, model_info: ModelInfo, } #[derive(Debug, Clone, Deserialize, Default)] pub struct OpenAIConfig { + pub name: Option, pub api_key: Option, pub organization_id: Option, - pub proxy: Option, - /// Set a timeout in seconds for connect to openai server - pub connect_timeout: Option, + pub extra: Option, } #[async_trait] impl Client for OpenAIClient { - fn get_config(&self) -> &SharedConfig { + fn config(&self) -> &SharedConfig { &self.global_config } - async fn send_message_inner(&self, content: &str) -> Result { - let builder = self.request_builder(content, false)?; + fn extra_config(&self) -> &Option { + &self.config.extra + } + + async fn send_message_inner(&self, client: &ReqwestClient, data: SendData) -> Result { + 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 OpenAIClient { + pub const NAME: &str = "openai"; + pub fn init(global_config: SharedConfig) -> Option> { let model_info = global_config.read().model_info.clone(); - if model_info.client != OpenAIClient::name() { - return None; - } - let local_config = { + let config = { if let ClientConfig::OpenAI(c) = &global_config.read().clients[model_info.index] { c.clone() } else { @@ -68,16 +69,18 @@ impl OpenAIClient { }; Some(Box::new(Self { global_config, - local_config, + config, model_info, })) } - pub fn name() -> &'static str { - "openai" + pub fn name(local_config: &OpenAIConfig) -> &str { + local_config.name.as_deref().unwrap_or(Self::NAME) } - pub fn list_models(_local_config: &OpenAIConfig, index: usize) -> Vec { + pub fn list_models(local_config: &OpenAIConfig, index: usize) -> Vec { + let client = Self::name(local_config); + [ ("gpt-3.5-turbo", 4096), ("gpt-3.5-turbo-16k", 16384), @@ -85,90 +88,38 @@ impl OpenAIClient { ("gpt-4-32k", 32768), ] .into_iter() - .map(|(name, max_tokens)| ModelInfo::new(Self::name(), name, max_tokens, index)) + .map(|(name, max_tokens)| ModelInfo::new(client, name, max_tokens, index)) .collect() } pub fn create_config() -> Result { - 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."))?; + 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")); - 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 { - 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("OPENAI_API_KEY") { - api_key.to_string() - } else { - bail!("Miss api_key") - }; + fn request_builder(&self, client: &ReqwestClient, data: SendData) -> Result { + let env_prefix = Self::name(&self.config).to_uppercase(); - let messages = self.global_config.read().build_messages(content)?; + 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 mut body = json!({ - "model": self.model_info.name, - "messages": messages, - }); + let body = openai_build_body(data, self.model_info.name.clone()); - 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 api_base = env::var("OPENAI_API_BASE") + let api_base = env::var(format!("{env_prefix}_API_BASE")) .ok() .unwrap_or_else(|| API_BASE.to_string()); + let url = format!("{api_base}/chat/completions"); let mut builder = client.post(url).bearer_auth(api_key).json(&body); - if let Some(organization_id) = &self.local_config.organization_id { + if let Some(organization_id) = &self.config.organization_id { builder = builder.header("OpenAI-Organization", organization_id); } @@ -219,3 +170,26 @@ pub(crate) async fn openai_send_message_streaming( Ok(()) } + +pub(crate) 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))); + } + + if stream { + body.as_object_mut() + .and_then(|m| m.insert("stream".into(), json!(true))); + } + body +} -- cgit v1.2.3