diff options
Diffstat (limited to 'src/client/mod.rs')
| -rw-r--r-- | src/client/mod.rs | 14 |
1 files changed, 13 insertions, 1 deletions
diff --git a/src/client/mod.rs b/src/client/mod.rs index 810fac2..ca37804 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -1,7 +1,9 @@ +pub mod azure_openai; pub mod localai; pub mod openai; use self::{ + azure_openai::{AzureOpenAIClient, AzureOpenAIConfig}, localai::LocalAIConfig, openai::{OpenAIClient, OpenAIConfig}, }; @@ -27,6 +29,8 @@ pub enum ClientConfig { OpenAI(OpenAIConfig), #[serde(rename = "localai")] LocalAI(LocalAIConfig), + #[serde(rename = "azure-openai")] + AzureOpenAI(AzureOpenAIConfig), } #[derive(Debug, Clone)] @@ -128,6 +132,7 @@ pub trait Client { 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!( @@ -139,7 +144,11 @@ pub fn init_client(config: SharedConfig) -> Result<Box<dyn Client>> { } pub fn all_clients() -> Vec<&'static str> { - vec![OpenAIClient::name(), LocalAIClient::name()] + vec![ + OpenAIClient::name(), + LocalAIClient::name(), + AzureOpenAIClient::name(), + ] } pub fn create_client_config(client: &str) -> Result<String> { @@ -147,6 +156,8 @@ pub fn create_client_config(client: &str) -> Result<String> { OpenAIClient::create_config() } else if client == LocalAIClient::name() { LocalAIClient::create_config() + } else if client == AzureOpenAIClient::name() { + AzureOpenAIClient::create_config() } else { bail!("Unknown client {}", &client) } @@ -160,6 +171,7 @@ pub fn list_models(config: &Config) -> Vec<ModelInfo> { .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() } |