From a491820ab5ff9bb58e62fd25d88f31ab53a81db9 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Thu, 6 Feb 2025 19:20:48 +0800
Subject: feat: support model alias (#1150)

---
 src/client/vertexai.rs | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

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

diff --git a/src/client/vertexai.rs b/src/client/vertexai.rs
index 19c8436..7cf5a1f 100644
--- a/src/client/vertexai.rs
+++ b/src/client/vertexai.rs
@@ -43,7 +43,7 @@ impl Client for VertexAIClient {
     ) -> Result<ChatCompletionsOutput> {
         prepare_gcloud_access_token(client, self.name(), &self.config.adc_file).await?;
         let model = self.model();
-        let model_category = ModelCategory::from_str(model.name())?;
+        let model_category = ModelCategory::from_str(model.real_name())?;
         let request_data = prepare_chat_completions(self, data, &model_category)?;
         let builder = self.request_builder(client, request_data);
         match model_category {
@@ -61,7 +61,7 @@ impl Client for VertexAIClient {
     ) -> Result<()> {
         prepare_gcloud_access_token(client, self.name(), &self.config.adc_file).await?;
         let model = self.model();
-        let model_category = ModelCategory::from_str(model.name())?;
+        let model_category = ModelCategory::from_str(model.real_name())?;
         let request_data = prepare_chat_completions(self, data, &model_category)?;
         let builder = self.request_builder(client, request_data);
         match model_category {
@@ -100,7 +100,7 @@ fn prepare_chat_completions(
 
     let base_url = format!("https://{location}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{location}/publishers");
 
-    let model_name = self_.model.name();
+    let model_name = self_.model.real_name();
 
     let url = match model_category {
         ModelCategory::Gemini => {
@@ -135,7 +135,7 @@ fn prepare_chat_completions(
         ModelCategory::Mistral => {
             let mut body = openai_build_chat_completions_body(data, &self_.model);
             if let Some(body_obj) = body.as_object_mut() {
-                body_obj["model"] = strip_model_version(self_.model.name()).into();
+                body_obj["model"] = strip_model_version(self_.model.real_name()).into();
             }
             body
         }
@@ -154,7 +154,7 @@ fn prepare_embeddings(self_: &VertexAIClient, data: &EmbeddingsData) -> Result<R
     let access_token = get_access_token(self_.name())?;
 
     let base_url = format!("https://{location}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{location}/publishers");
-    let url = format!("{base_url}/google/models/{}:predict", self_.model.name());
+    let url = format!("{base_url}/google/models/{}:predict", self_.model.real_name());
 
     let instances: Vec<_> = data.texts.iter().map(|v| json!({"content": v})).collect();
 
-- 
cgit v1.2.3