From abc588daac6053ec2edbdcde3f5a2dc5eb7d50b8 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Fri, 21 Jun 2024 06:00:26 +0800
Subject: feat: support rerank (#620)

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

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

diff --git a/src/client/model.rs b/src/client/model.rs
index 56421bf..ebf1264 100644
--- a/src/client/model.rs
+++ b/src/client/model.rs
@@ -1,5 +1,5 @@
 use super::{
-    list_chat_models, list_embedding_models,
+    list_chat_models, list_embedding_models, list_rerank_models,
     message::{Message, MessageContent},
     EmbeddingsData,
 };
@@ -46,14 +46,21 @@ impl Model {
     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!("Invalid model '{model_id}'"),
+            None => bail!("Invalid 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!("Invalid model '{model_id}'"),
+            None => bail!("Invalid embedding model '{model_id}'"),
+        }
+    }
+
+    pub fn retrieve_rerank(config: &Config, model_id: &str) -> Result<Self> {
+        match Self::find(&list_rerank_models(config), model_id) {
+            Some(v) => Ok(v),
+            None => bail!("Invalid rerank model '{model_id}'"),
         }
     }
 
-- 
cgit v1.2.3