From 3a3388375be05758d5f5574cce1d0f8eb7bdd604 Mon Sep 17 00:00:00 2001 From: sigoden Date: Wed, 4 Dec 2024 21:03:59 +0800 Subject: refactor: improve retrieve model (#1036) - check the model type while retrieve model - select chat/reranker model even if it is missed in client models - find predefined-models for openai-compatible client with startsWith - remove client::ApiType --- src/client/bedrock.rs | 4 +- src/client/common.rs | 38 +++------------ src/client/ernie.rs | 8 ++-- src/client/macros.rs | 45 +++++++++++------- src/client/model.rs | 124 +++++++++++++++++++++++++++++++++---------------- src/client/vertexai.rs | 6 +-- 6 files changed, 126 insertions(+), 99 deletions(-) (limited to 'src/client') diff --git a/src/client/bedrock.rs b/src/client/bedrock.rs index 081a65b..7cd289c 100644 --- a/src/client/bedrock.rs +++ b/src/client/bedrock.rs @@ -67,7 +67,7 @@ impl BedrockClient { let body = build_chat_completions_body(data, &self.model)?; let mut request_data = RequestData::new("", body); - self.patch_request_data(&mut request_data, ApiType::ChatCompletions); + self.patch_request_data(&mut request_data); let RequestData { url: _, headers, @@ -118,7 +118,7 @@ impl BedrockClient { }); let mut request_data = RequestData::new("", body); - self.patch_request_data(&mut request_data, ApiType::Embeddings); + self.patch_request_data(&mut request_data); let RequestData { url: _, headers, diff --git a/src/client/common.rs b/src/client/common.rs index 9c516c8..e0ec861 100644 --- a/src/client/common.rs +++ b/src/client/common.rs @@ -19,7 +19,7 @@ use tokio::sync::mpsc::unbounded_channel; const MODELS_YAML: &str = include_str!("../../models.yaml"); lazy_static::lazy_static! { - pub static ref ALL_MODELS: Vec = serde_yaml::from_str(MODELS_YAML).unwrap(); + pub static ref ALL_PREDEFINED_MODELS: Vec = serde_yaml::from_str(MODELS_YAML).unwrap(); static ref ESCAPE_SLASH_RE: Regex = Regex::new(r"(? RequestBuilder { - self.patch_request_data(&mut request_data, api_type); + self.patch_request_data(&mut request_data); request_data.into_builder(client) } - fn patch_request_data(&self, request_data: &mut RequestData, api_type: ApiType) { + fn patch_request_data(&self, request_data: &mut RequestData) { + let model_type = self.model().model_type(); let map = std::env::var(get_env_name(&format!( "patch_{}_{}", self.model().client_name(), - api_type.name(), + model_type.api_name(), ))) .ok() .and_then(|v| serde_json::from_str(&v).ok()) .or_else(|| { self.patch_config() - .and_then(|v| api_type.extract_patch(v)) + .and_then(|v| model_type.extract_patch(v)) .cloned() }); let map = match map { @@ -200,30 +200,6 @@ pub struct RequestPatch { pub type ApiPatch = IndexMap; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ApiType { - ChatCompletions, - Embeddings, - Rerank, -} - -impl ApiType { - pub fn name(&self) -> &str { - match self { - ApiType::ChatCompletions => "chat_completions", - ApiType::Embeddings => "embeddings", - ApiType::Rerank => "rerank", - } - } - pub fn extract_patch<'a>(&self, patch: &'a RequestPatch) -> Option<&'a ApiPatch> { - match self { - ApiType::ChatCompletions => patch.chat_completions.as_ref(), - ApiType::Embeddings => patch.embeddings.as_ref(), - ApiType::Rerank => patch.rerank.as_ref(), - } - } -} - pub struct RequestData { pub url: String, pub headers: IndexMap, @@ -383,7 +359,7 @@ pub fn create_openai_compatible_client_config(client: &str) -> Result Result { prepare_access_token(self, client).await?; let request_data = prepare_chat_completions(self, data)?; - let builder = self.request_builder(client, request_data, ApiType::ChatCompletions); + let builder = self.request_builder(client, request_data); chat_completions(builder, &self.model).await } @@ -53,7 +53,7 @@ impl Client for ErnieClient { ) -> Result<()> { prepare_access_token(self, client).await?; let request_data = prepare_chat_completions(self, data)?; - let builder = self.request_builder(client, request_data, ApiType::ChatCompletions); + let builder = self.request_builder(client, request_data); chat_completions_streaming(builder, handler, &self.model).await } @@ -64,7 +64,7 @@ impl Client for ErnieClient { ) -> Result { prepare_access_token(self, client).await?; let request_data = prepare_embeddings(self, data)?; - let builder = self.request_builder(client, request_data, ApiType::Embeddings); + let builder = self.request_builder(client, request_data); embeddings(builder, &self.model).await } @@ -75,7 +75,7 @@ impl Client for ErnieClient { ) -> Result { prepare_access_token(self, client).await?; let request_data = prepare_rerank(self, data)?; - let builder = self.request_builder(client, request_data, ApiType::Rerank); + let builder = self.request_builder(client, request_data); rerank(builder, &self.model).await } } diff --git a/src/client/macros.rs b/src/client/macros.rs index 0543cd6..4f52044 100644 --- a/src/client/macros.rs +++ b/src/client/macros.rs @@ -52,9 +52,10 @@ macro_rules! register_client { pub fn list_models(local_config: &$config) -> Vec { let client_name = Self::name(local_config); if local_config.models.is_empty() { - if let Some(models) = $crate::client::ALL_MODELS.iter().find(|v| { + if let Some(models) = $crate::client::ALL_PREDEFINED_MODELS.iter().find(|v| { v.platform == $name || - ($name == OpenAICompatibleClient::NAME && local_config.name.as_deref() == Some(&v.platform)) + ($name == OpenAICompatibleClient::NAME + && local_config.name.as_ref().map(|name| name.starts_with(&v.platform)).unwrap_or_default()) }) { return Model::from_config(client_name, &models.models); } @@ -98,32 +99,40 @@ macro_rules! register_client { anyhow::bail!("Unknown client '{}'", client) } - static ALL_CLIENT_MODELS: std::sync::OnceLock> = std::sync::OnceLock::new(); + static ALL_CLIENT_NAMES: std::sync::OnceLock> = std::sync::OnceLock::new(); - pub fn list_models(config: &$crate::config::Config) -> Vec<&'static $crate::client::Model> { - let models = ALL_CLIENT_MODELS.get_or_init(|| { + pub fn list_client_names(config: &$crate::config::Config) -> Vec<&'static String> { + let names = ALL_CLIENT_NAMES.get_or_init(|| { config .clients .iter() .flat_map(|v| match v { - $(ClientConfig::$config(c) => $client::list_models(c),)+ + $(ClientConfig::$config(c) => vec![$client::name(c).to_string()],)+ ClientConfig::Unknown => vec![], }) .collect() }); - models.iter().collect() + names.iter().collect() } - pub fn list_chat_models(config: &$crate::config::Config) -> Vec<&'static $crate::client::Model> { - list_models(config).into_iter().filter(|v| v.model_type() == "chat").collect() - } + static ALL_MODELS: std::sync::OnceLock> = std::sync::OnceLock::new(); - pub fn list_embedding_models(config: &$crate::config::Config) -> Vec<&'static $crate::client::Model> { - list_models(config).into_iter().filter(|v| v.model_type() == "embedding").collect() + pub fn list_all_models(config: &$crate::config::Config) -> Vec<&'static $crate::client::Model> { + let models = ALL_MODELS.get_or_init(|| { + config + .clients + .iter() + .flat_map(|v| match v { + $(ClientConfig::$config(c) => $client::list_models(c),)+ + ClientConfig::Unknown => vec![], + }) + .collect() + }); + models.iter().collect() } - pub fn list_reranker_models(config: &$crate::config::Config) -> Vec<&'static $crate::client::Model> { - list_models(config).into_iter().filter(|v| v.model_type() == "reranker").collect() + pub fn list_models(config: &$crate::config::Config, model_type: $crate::client::ModelType) -> Vec<&'static $crate::client::Model> { + list_all_models(config).into_iter().filter(|v| v.model_type() == model_type).collect() } }; } @@ -175,7 +184,7 @@ macro_rules! impl_client_trait { data: $crate::client::ChatCompletionsData, ) -> anyhow::Result<$crate::client::ChatCompletionsOutput> { let request_data = $prepare_chat_completions(self, data)?; - let builder = self.request_builder(client, request_data, ApiType::ChatCompletions); + let builder = self.request_builder(client, request_data); $chat_completions(builder, self.model()).await } @@ -186,7 +195,7 @@ macro_rules! impl_client_trait { data: $crate::client::ChatCompletionsData, ) -> Result<()> { let request_data = $prepare_chat_completions(self, data)?; - let builder = self.request_builder(client, request_data, ApiType::ChatCompletions); + let builder = self.request_builder(client, request_data); $chat_completions_streaming(builder, handler, self.model()).await } @@ -196,7 +205,7 @@ macro_rules! impl_client_trait { data: &$crate::client::EmbeddingsData, ) -> Result<$crate::client::EmbeddingsOutput> { let request_data = $prepare_embeddings(self, data)?; - let builder = self.request_builder(client, request_data, ApiType::Embeddings); + let builder = self.request_builder(client, request_data); $embeddings(builder, self.model()).await } @@ -206,7 +215,7 @@ macro_rules! impl_client_trait { data: &$crate::client::RerankData, ) -> Result<$crate::client::RerankOutput> { let request_data = $prepare_rerank(self, data)?; - let builder = self.request_builder(client, request_data, ApiType::Rerank); + let builder = self.request_builder(client, request_data); $rerank(builder, self.model()).await } } diff --git a/src/client/model.rs b/src/client/model.rs index 5d496d2..47f6598 100644 --- a/src/client/model.rs +++ b/src/client/model.rs @@ -1,7 +1,7 @@ use super::{ - list_chat_models, list_embedding_models, list_reranker_models, + list_all_models, list_client_names, message::{Message, MessageContent, MessageContentPart}, - MessageContentToolCalls, + ApiPatch, MessageContentToolCalls, RequestPatch, }; use crate::config::Config; @@ -9,6 +9,7 @@ use crate::utils::{estimate_token_length, format_option_value}; use anyhow::{bail, Result}; use serde::{Deserialize, Serialize}; +use std::fmt::Display; const PER_MESSAGES_TOKENS: usize = 5; const BASIS_TOKENS: usize = 2; @@ -43,29 +44,8 @@ impl Model { .collect() } - pub fn retrieve_chat(config: &Config, model_id: &str) -> Result { - match Self::find(&list_chat_models(config), model_id) { - Some(v) => Ok(v), - None => bail!("Unknown chat model '{model_id}'"), - } - } - - pub fn retrieve_embedding(config: &Config, model_id: &str) -> Result { - match Self::find(&list_embedding_models(config), model_id) { - Some(v) => Ok(v), - None => bail!("Unknown embedding model '{model_id}'"), - } - } - - pub fn retrieve_reranker(config: &Config, model_id: &str) -> Result { - match Self::find(&list_reranker_models(config), model_id) { - Some(v) => Ok(v), - None => bail!("Unknown reranker model '{model_id}'"), - } - } - - pub fn find(models: &[&Self], model_id: &str) -> Option { - let mut model = None; + pub fn retrieve_model(config: &Config, model_id: &str, model_type: ModelType) -> Result { + let models = list_all_models(config); let (client_name, model_name) = match model_id.split_once(':') { Some((client_name, model_name)) => { if model_name.is_empty() { @@ -78,21 +58,33 @@ impl Model { }; match model_name { Some(model_name) => { - if let Some(found) = models.iter().find(|v| v.id() == model_id) { - model = Some((*found).clone()); - } else if let Some(found) = models.iter().find(|v| v.client_name == client_name) { - let mut found = (*found).clone(); - found.data.name = model_name.to_string(); - model = Some(found) + if let Some(model) = models.iter().find(|v| v.id() == model_id) { + if model.model_type() == model_type { + return Ok((*model).clone()); + } else { + bail!("Model '{model_id}' is not a {model_type} model") + } + } + if list_client_names(config) + .into_iter() + .any(|v| *v == client_name) + && model_type.can_create_from_name() + { + let mut new_model = Self::new(client_name, model_name); + new_model.data.model_type = model_type.to_string(); + return Ok(new_model); } } None => { - if let Some(found) = models.iter().find(|v| v.client_name == client_name) { - model = Some((*found).clone()); + if let Some(found) = models + .iter() + .find(|v| v.client_name == client_name && v.model_type() == model_type) + { + return Ok((*found).clone()); } } - } - model + }; + bail!("Unknown {model_type} model '{model_id}'") } pub fn id(&self) -> String { @@ -111,8 +103,14 @@ impl Model { &self.data.name } - pub fn model_type(&self) -> &str { - &self.data.model_type + pub fn model_type(&self) -> ModelType { + if self.data.model_type.starts_with("embed") { + ModelType::Embedding + } else if self.data.model_type.starts_with("rerank") { + ModelType::Reranker + } else { + ModelType::Chat + } } pub fn data(&self) -> &ModelData { @@ -125,7 +123,7 @@ impl Model { pub fn description(&self) -> String { match self.model_type() { - "chat" => { + ModelType::Chat => { let ModelData { max_input_tokens, max_output_tokens, @@ -156,7 +154,7 @@ impl Model { max_input_tokens, max_output_tokens, input_price, output_price, capabilities ) } - "embedding" => { + ModelType::Embedding => { let ModelData { input_price, max_tokens_per_chunk, @@ -168,7 +166,7 @@ impl Model { let price = format_option_value(input_price); format!("max-tokens:{max_tokens};max-batch:{max_batch};price:{price}") } - _ => String::new(), + ModelType::Reranker => String::new(), } } @@ -310,13 +308,14 @@ impl ModelData { pub fn new(name: &str) -> Self { Self { name: name.to_string(), + model_type: default_model_type(), ..Default::default() } } } #[derive(Debug, Clone, Deserialize)] -pub struct BuiltinModels { +pub struct PredefinedModels { pub platform: String, pub models: Vec, } @@ -324,3 +323,46 @@ pub struct BuiltinModels { fn default_model_type() -> String { "chat".into() } + +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum ModelType { + Chat, + Embedding, + Reranker, +} + +impl Display for ModelType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ModelType::Chat => write!(f, "chat"), + ModelType::Embedding => write!(f, "embedding"), + ModelType::Reranker => write!(f, "reranker"), + } + } +} + +impl ModelType { + pub fn can_create_from_name(self) -> bool { + match self { + ModelType::Chat => true, + ModelType::Embedding => false, + ModelType::Reranker => true, + } + } + + pub fn api_name(self) -> &'static str { + match self { + ModelType::Chat => "chat_completions", + ModelType::Embedding => "embeddings", + ModelType::Reranker => "rerank", + } + } + + pub fn extract_patch(self, patch: &RequestPatch) -> Option<&ApiPatch> { + match self { + ModelType::Chat => patch.chat_completions.as_ref(), + ModelType::Embedding => patch.embeddings.as_ref(), + ModelType::Reranker => patch.rerank.as_ref(), + } + } +} diff --git a/src/client/vertexai.rs b/src/client/vertexai.rs index 07ff2a4..7b73164 100644 --- a/src/client/vertexai.rs +++ b/src/client/vertexai.rs @@ -45,7 +45,7 @@ impl Client for VertexAIClient { let model = self.model(); let model_category = ModelCategory::from_str(model.name())?; let request_data = prepare_chat_completions(self, data, &model_category)?; - let builder = self.request_builder(client, request_data, ApiType::ChatCompletions); + let builder = self.request_builder(client, request_data); match model_category { ModelCategory::Gemini => gemini_chat_completions(builder, model).await, ModelCategory::Claude => claude_chat_completions(builder, model).await, @@ -63,7 +63,7 @@ impl Client for VertexAIClient { let model = self.model(); let model_category = ModelCategory::from_str(model.name())?; let request_data = prepare_chat_completions(self, data, &model_category)?; - let builder = self.request_builder(client, request_data, ApiType::ChatCompletions); + let builder = self.request_builder(client, request_data); match model_category { ModelCategory::Gemini => { gemini_chat_completions_streaming(builder, handler, model).await @@ -84,7 +84,7 @@ impl Client for VertexAIClient { ) -> Result>> { prepare_gcloud_access_token(client, self.name(), &self.config.adc_file).await?; let request_data = prepare_embeddings(self, data)?; - let builder = self.request_builder(client, request_data, ApiType::Embeddings); + let builder = self.request_builder(client, request_data); embeddings(builder, self.model()).await } } -- cgit v1.2.3