From 9c6c9f10a27d0993636b453f39d8934c95c5c2b2 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Tue, 23 Apr 2024 18:14:47 +0800
Subject: feat: builtin models can be overwrited by models config (#429)

---
 src/client/model.rs | 37 ++++++++++++++++++++++++-------------
 1 file changed, 24 insertions(+), 13 deletions(-)

(limited to 'src/client/model.rs')

diff --git a/src/client/model.rs b/src/client/model.rs
index 53d1834..b97c244 100644
--- a/src/client/model.rs
+++ b/src/client/model.rs
@@ -36,6 +36,30 @@ impl Model {
         }
     }
 
+    pub fn from_config(client_name: &str, models: &[ModelConfig]) -> Vec<Self> {
+        models
+            .iter()
+            .map(|v| {
+                Model::new(client_name, &v.name)
+                    .set_capabilities(v.capabilities)
+                    .set_max_input_tokens(v.max_input_tokens)
+                    .set_max_output_tokens(v.max_output_tokens)
+                    .set_extra_fields(v.extra_fields.clone())
+            })
+            .collect()
+    }
+
+    pub fn from_static(client_name: &str, models: &[(&str, usize, &str)]) -> Vec<Self> {
+        models
+            .iter()
+            .map(|(name, max_input_tokens, capabilities)| {
+                Model::new(client_name, name)
+                    .set_capabilities((*capabilities).into())
+                    .set_max_input_tokens(Some(*max_input_tokens))
+            })
+            .collect()
+    }
+
     pub fn find(models: &[Self], value: &str) -> Option<Self> {
         let mut model = None;
         let (client_name, model_name) = match value.split_once(':') {
@@ -156,19 +180,6 @@ impl Model {
     }
 }
 
-pub fn convert_models(client_name: &str, models: &[ModelConfig]) -> Vec<Model> {
-    models
-        .iter()
-        .map(|v| {
-            Model::new(client_name, &v.name)
-                .set_capabilities(v.capabilities)
-                .set_max_input_tokens(v.max_input_tokens)
-                .set_max_output_tokens(v.max_output_tokens)
-                .set_extra_fields(v.extra_fields.clone())
-        })
-        .collect()
-}
-
 #[derive(Debug, Clone, Deserialize)]
 pub struct ModelConfig {
     pub name: String,
-- 
cgit v1.2.3