From 602494b650a59b73b4548b9282375c99b471e11b Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Mon, 29 Apr 2024 11:57:02 +0800
Subject: refactor: merge config models, update client models (#460)

---
 src/client/common.rs | 16 ++++++----------
 src/client/model.rs  | 13 +++++++++++--
 2 files changed, 17 insertions(+), 12 deletions(-)

(limited to 'src/client')

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)
-- 
cgit v1.2.3