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

1 files changed, 83 insertions, 41 deletions

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<Self> {
-        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<Self> {
-        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<Self> {
-        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<Self> {
-        let mut model = None;
+    pub fn retrieve_model(config: &Config, model_id: &str, model_type: ModelType) -> Result<Self> {
+        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<ModelData>,
 }
@@ -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(),
+        }
+    }
+}