From c877ae7fc6d123c68263adc99f1832f4ca5e98f8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 14 Nov 2024 14:35:09 -0800 Subject: [PATCH] fix use fine tuned vertex embedding models --- litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py b/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py index 0f95b222c..74bab0b26 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/common_utils.py @@ -89,6 +89,9 @@ def _get_vertex_url( elif mode == "embedding": endpoint = "predict" url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}" + if model.isdigit(): + # https://us-central1-aiplatform.googleapis.com/v1/projects/$PROJECT_ID/locations/us-central1/endpoints/$ENDPOINT_ID:predict + url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}:{endpoint}" if not url or not endpoint: raise ValueError(f"Unable to get vertex url/endpoint for mode: {mode}")