diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/client/common.rs | 16 | ||||
| -rw-r--r-- | src/client/model.rs | 13 |
2 files changed, 17 insertions, 12 deletions
diff --git a/src/client/common.rs b/src/client/common.rs index 842741c..aa243a7 100644 --- a/src/client/common.rs +++ b/src/client/common.rs @@ -89,19 +89,15 @@ macro_rules! register_client { pub fn list_models(local_config: &$config) -> Vec<Model> { let client_name = Self::name(local_config); - if local_config.models.is_empty() { - for model in $crate::client::CLIENT_MODELS.iter() { - match model { - $crate::client::ClientModel::$config { models } => { - return Model::from_config(client_name, models); - } - _ => {} + for model in $crate::client::CLIENT_MODELS.iter() { + match model { + $crate::client::ClientModel::$config { models } => { + return Model::from_config(client_name, &local_config.models, models); } + _ => {} } - vec![] - } else { - Model::from_config(client_name, &local_config.models) } + vec![] } pub fn name(config: &$config) -> &str { diff --git a/src/client/model.rs b/src/client/model.rs index 459d94e..f3fe97c 100644 --- a/src/client/model.rs +++ b/src/client/model.rs @@ -3,6 +3,7 @@ use super::message::{Message, MessageContent}; use crate::utils::count_tokens; use anyhow::{bail, Result}; +use indexmap::IndexMap; use serde::Deserialize; const PER_MESSAGES_TOKENS: usize = 5; @@ -36,9 +37,17 @@ impl Model { } } - pub fn from_config(client_name: &str, models: &[ModelConfig]) -> Vec<Self> { + pub fn from_config( + client_name: &str, + config_models: &[ModelConfig], + client_models: &[ModelConfig], + ) -> Vec<Self> { + let mut models = IndexMap::new(); + for model in client_models.iter().chain(config_models.iter()) { + models.insert(&model.name, model); + } models - .iter() + .values() .map(|v| { Model::new(client_name, &v.name) .set_max_input_tokens(v.max_input_tokens) |