update

sigoden · Dec 4, 2024 · 931b717 · 931b717
1 parent 236a499
commit 931b717
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/src/client/common.rs b/src/client/common.rs
@@ -150,16 +150,17 @@ pub trait Client: Sync + Send {
     }
 
     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(),
-            self.model().model_type().api_name(),
+            model_type.api_name(),
         )))
         .ok()
         .and_then(|v| serde_json::from_str(&v).ok())
         .or_else(|| {
             self.patch_config()
-                .and_then(|v| self.model().model_type().extract_patch(v))
+                .and_then(|v| model_type.extract_patch(v))
                 .cloned()
         });
         let map = match map {

diff --git a/src/client/model.rs b/src/client/model.rs
@@ -104,10 +104,12 @@ impl Model {
     }
 
     pub fn model_type(&self) -> ModelType {
-        match self.data.model_type.as_str() {
-            "embed" | "embedding" => ModelType::Embedding,
-            "rerank" | "reranker" => ModelType::Reranker,
-            _ => ModelType::Chat,
+        if self.data.model_type.starts_with("embed") {
+            return ModelType::Embedding;
+        } else if self.data.model_type.starts_with("rerank") {
+            return ModelType::Reranker;
+        } else {
+            ModelType::Chat
         }
     }