From bf1308e86bdcb3e86fb5ef2b976418dabe910c72 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 29 Oct 2025 06:21:35 +0530 Subject: [PATCH 01/23] Support for Custom Vertex AI Models via PSC Endpoint with api_base (#15953) * Support for Custom Vertex AI Models via PSC Endpoint with api_base * Add docs related psc * remove not needed files * remove print statemnt * fix mypy errors --- docs/my-website/docs/providers/vertex.md | 47 ++++ litellm/llms/vertex_ai/batches/handler.py | 8 + litellm/llms/vertex_ai/common_utils.py | 10 +- .../vertex_ai_context_caching.py | 4 + .../vertex_embeddings/transformation.py | 3 + litellm/llms/vertex_ai/vertex_llm_base.py | 48 +++- .../vertex_ai/vertex_model_garden/main.py | 4 + .../test_vertex_ai_psc_endpoint_support.py | 258 ++++++++++++++++++ 8 files changed, 378 insertions(+), 4 deletions(-) create mode 100644 tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 70babea3814..8e333b69ef7 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -1604,6 +1604,53 @@ litellm.vertex_location = "us-central1 # Your Location | gemini-2.5-flash-preview-09-2025 | `completion('gemini-2.5-flash-preview-09-2025', messages)`, `completion('vertex_ai/gemini-2.5-flash-preview-09-2025', messages)` | | gemini-2.5-flash-lite-preview-09-2025 | `completion('gemini-2.5-flash-lite-preview-09-2025', messages)`, `completion('vertex_ai/gemini-2.5-flash-lite-preview-09-2025', messages)` | +## Private Service Connect (PSC) Endpoints + +LiteLLM supports Vertex AI models deployed to Private Service Connect (PSC) endpoints, allowing you to use custom `api_base` URLs for private deployments. + +### Usage + +```python +from litellm import completion + +# Use PSC endpoint with custom api_base +response = completion( + model="vertex_ai/1234567890", # Numeric endpoint ID + messages=[{"role": "user", "content": "Hello!"}], + api_base="http://10.96.32.8", # Your PSC endpoint + vertex_project="my-project-id", + vertex_location="us-central1" +) +``` + +**Key Features:** +- Supports both numeric endpoint IDs and custom model names +- Works with both completion and embedding endpoints +- Automatically constructs full PSC URL: `{api_base}/v1/projects/{project}/locations/{location}/endpoints/{model}:{endpoint}` +- Compatible with streaming requests + +### Configuration + +Add PSC endpoints to your `config.yaml`: + +```yaml +model_list: + - model_name: psc-gemini + litellm_params: + model: vertex_ai/1234567890 # Numeric endpoint ID + api_base: "http://10.96.32.8" # Your PSC endpoint + vertex_project: "my-project-id" + vertex_location: "us-central1" + vertex_credentials: "/path/to/service_account.json" + - model_name: psc-embedding + litellm_params: + model: vertex_ai/text-embedding-004 + api_base: "http://10.96.32.8" # Your PSC endpoint + vertex_project: "my-project-id" + vertex_location: "us-central1" + vertex_credentials: "/path/to/service_account.json" +``` + ## Fine-tuned Models You can call fine-tuned Vertex AI Gemini models through LiteLLM diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 864cc190312..edae91ff9a3 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -61,6 +61,10 @@ class VertexAIBatchPrediction(VertexLLM): stream=None, auth_header=None, url=default_api_base, + model=None, + vertex_project=vertex_project or project_id, + vertex_location=vertex_location or "us-central1", + vertex_api_version="v1", ) headers = { @@ -166,6 +170,10 @@ class VertexAIBatchPrediction(VertexLLM): stream=None, auth_header=None, url=default_api_base, + model=None, + vertex_project=vertex_project or project_id, + vertex_location=vertex_location or "us-central1", + vertex_api_version="v1", ) headers = { diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index dc6a3170afe..aaee922a3f0 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -60,6 +60,9 @@ def get_vertex_ai_model_route( >>> get_vertex_ai_model_route("openai/gpt-oss-120b") VertexAIModelRoute.MODEL_GARDEN + + >>> get_vertex_ai_model_route("1234567890", {"api_base": "http://10.96.32.8"}) + VertexAIModelRoute.GEMINI # Numeric endpoints with api_base use HTTP path """ from litellm.llms.vertex_ai.vertex_ai_partner_models.main import ( VertexAIPartnerModels, @@ -69,7 +72,12 @@ def get_vertex_ai_model_route( if litellm_params and litellm_params.get("base_model") is not None: if "gemini" in litellm_params["base_model"]: return VertexAIModelRoute.GEMINI - + + # Check if numeric endpoint ID with custom api_base (PSC endpoint) + # Route to GEMINI (HTTP path) to support PSC endpoints properly + if model.isdigit() and litellm_params and litellm_params.get("api_base"): + return VertexAIModelRoute.GEMINI + # Check for partner models (llama, mistral, claude, etc.) if VertexAIPartnerModels.is_vertex_partner_model(model=model): return VertexAIModelRoute.PARTNER_MODELS diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 26be4d3c2b8..cff1bebceb9 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -85,6 +85,10 @@ class ContextCachingEndpoints(VertexBase): stream=None, auth_header=auth_header, url=url, + model=None, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_api_version="v1beta1" if custom_llm_provider == "vertex_ai_beta" else "v1", ) def check_cache( diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index 97af558041d..caaf00e199e 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -167,6 +167,9 @@ class VertexAITextEmbeddingConfig(BaseModel): vertex_request["parameters"] = TextEmbeddingFineTunedParameters( **optional_params ) + # Remove 'shared_session' from parameters if present + if vertex_request["parameters"] is not None and "shared_session" in vertex_request["parameters"]: + del vertex_request["parameters"]["shared_session"] # type: ignore[typeddict-item] return vertex_request diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 9ddbc461a70..a5c44617fab 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -241,6 +241,9 @@ class VertexBase: auth_header=None, url=default_api_base, model=model, + vertex_project=vertex_project or project_id, + vertex_location=vertex_location or "us-central1", + vertex_api_version="v1", # Partner models typically use v1 ) return api_base @@ -289,9 +292,18 @@ class VertexBase: auth_header: Optional[str], url: str, model: Optional[str] = None, + vertex_project: Optional[str] = None, + vertex_location: Optional[str] = None, + vertex_api_version: Optional[Literal["v1", "v1beta1"]] = None, ) -> Tuple[Optional[str], str]: """ for cloudflare ai gateway - https://github.com/BerriAI/litellm/issues/4317 + + Handles custom api_base for: + 1. Gemini (Google AI Studio) - constructs /models/{model}:{endpoint} + 2. Vertex AI with standard proxies - constructs {api_base}:{endpoint} + 3. Vertex AI with PSC endpoints - constructs full path structure + {api_base}/v1/projects/{project}/locations/{location}/endpoints/{model}:{endpoint} ## Returns - (auth_header, url) - Tuple[Optional[str], str] @@ -311,8 +323,34 @@ class VertexBase: if gemini_api_key is not None: auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment] else: - url = "{}:{}".format(api_base, endpoint) - + # For Vertex AI + # Check if this is a PSC endpoint or custom deployment + # PSC/custom endpoints need the full path structure + if vertex_project and vertex_location and model: + # Check if model is numeric (endpoint ID) or if api_base doesn't contain googleapis.com + # These are indicators of PSC/custom endpoints + is_psc_or_custom = ( + "googleapis.com" not in api_base.lower() or model.isdigit() + ) + + if is_psc_or_custom: + # Construct full PSC/custom endpoint URL + # Format: {api_base}/v1/projects/{project}/locations/{location}/endpoints/{model}:{endpoint} + version = vertex_api_version or "v1" + url = "{}/{}/projects/{}/locations/{}/endpoints/{}:{}".format( + api_base.rstrip("/"), + version, + vertex_project, + vertex_location, + model, + endpoint, + ) + else: + # Standard proxy - just append endpoint + url = "{}:{}".format(api_base, endpoint) + else: + # Fallback to simple format if we don't have all parameters + url = "{}:{}".format(api_base, endpoint) if stream is True: url = url + "?alt=sse" return auth_header, url @@ -339,6 +377,7 @@ class VertexBase: Returns token, url """ + version: Optional[Literal["v1beta1", "v1"]] = None if custom_llm_provider == "gemini": url, endpoint = _get_gemini_url( mode=mode, @@ -354,7 +393,7 @@ class VertexBase: ) ### SET RUNTIME ENDPOINT ### - version: Literal["v1beta1", "v1"] = ( + version = ( "v1beta1" if should_use_v1beta1_features is True else "v1" ) url, endpoint = _get_vertex_url( @@ -375,6 +414,9 @@ class VertexBase: stream=stream, url=url, model=model, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_api_version=version, ) def _handle_reauthentication( diff --git a/litellm/llms/vertex_ai/vertex_model_garden/main.py b/litellm/llms/vertex_ai/vertex_model_garden/main.py index 1c57096734b..225e75a5add 100644 --- a/litellm/llms/vertex_ai/vertex_model_garden/main.py +++ b/litellm/llms/vertex_ai/vertex_model_garden/main.py @@ -123,6 +123,10 @@ class VertexAIModelGardenModels(VertexBase): stream=stream, auth_header=None, url=default_api_base, + model=model, + vertex_project=vertex_project or project_id, + vertex_location=vertex_location or "us-central1", + vertex_api_version="v1beta1", ) model = "" return openai_like_chat_completions.completion( diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py new file mode 100644 index 00000000000..46f365094c0 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py @@ -0,0 +1,258 @@ +""" +Unit tests for Vertex AI Private Service Connect (PSC) endpoint support + +Tests that LiteLLM properly constructs URLs when using custom api_base +for PSC endpoints. +""" + +import pytest +import sys +import os + +# Add the litellm package to the path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../../..")) + +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + +class TestVertexAIPSCEndpointSupport: + """Test cases for PSC endpoint URL construction""" + + def test_psc_endpoint_url_construction_basic(self): + """Test basic PSC endpoint URL construction for predict endpoint""" + vertex_base = VertexBase() + psc_api_base = "http://10.96.32.8" + endpoint_id = "1234567890" + project_id = "test-project" + location = "us-central1" + + auth_header, url = vertex_base._check_custom_proxy( + api_base=psc_api_base, + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="predict", + stream=False, + auth_header="test-token", + url="", # This will be replaced + model=endpoint_id, + vertex_project=project_id, + vertex_location=location, + vertex_api_version="v1", + ) + + expected_url = f"{psc_api_base}/v1/projects/{project_id}/locations/{location}/endpoints/{endpoint_id}:predict" + assert ( + url == expected_url + ), f"Expected {expected_url}, but got {url}" + + def test_psc_endpoint_url_construction_with_streaming(self): + """Test PSC endpoint URL construction with streaming enabled""" + vertex_base = VertexBase() + psc_api_base = "http://10.96.32.8" + endpoint_id = "1234567890" + project_id = "test-project" + location = "us-central1" + + auth_header, url = vertex_base._check_custom_proxy( + api_base=psc_api_base, + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="streamGenerateContent", + stream=True, + auth_header="test-token", + url="", + model=endpoint_id, + vertex_project=project_id, + vertex_location=location, + vertex_api_version="v1", + ) + + expected_url = f"{psc_api_base}/v1/projects/{project_id}/locations/{location}/endpoints/{endpoint_id}:streamGenerateContent?alt=sse" + assert ( + url == expected_url + ), f"Expected {expected_url}, but got {url}" + + def test_psc_endpoint_url_construction_v1beta1(self): + """Test PSC endpoint URL construction with v1beta1 API version""" + vertex_base = VertexBase() + psc_api_base = "http://10.96.32.8" + endpoint_id = "1234567890" + project_id = "test-project" + location = "us-central1" + + auth_header, url = vertex_base._check_custom_proxy( + api_base=psc_api_base, + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="predict", + stream=False, + auth_header="test-token", + url="", + model=endpoint_id, + vertex_project=project_id, + vertex_location=location, + vertex_api_version="v1beta1", + ) + + expected_url = f"{psc_api_base}/v1beta1/projects/{project_id}/locations/{location}/endpoints/{endpoint_id}:predict" + assert ( + url == expected_url + ), f"Expected {expected_url}, but got {url}" + + def test_psc_endpoint_url_with_https(self): + """Test PSC endpoint URL construction with HTTPS""" + vertex_base = VertexBase() + psc_api_base = "https://10.96.32.8" + endpoint_id = "1234567890" + project_id = "test-project" + location = "us-central1" + + auth_header, url = vertex_base._check_custom_proxy( + api_base=psc_api_base, + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="predict", + stream=False, + auth_header="test-token", + url="", + model=endpoint_id, + vertex_project=project_id, + vertex_location=location, + vertex_api_version="v1", + ) + + expected_url = f"{psc_api_base}/v1/projects/{project_id}/locations/{location}/endpoints/{endpoint_id}:predict" + assert ( + url == expected_url + ), f"Expected {expected_url}, but got {url}" + + def test_psc_endpoint_with_trailing_slash(self): + """Test that trailing slashes in api_base are handled correctly""" + vertex_base = VertexBase() + psc_api_base = "http://10.96.32.8/" + endpoint_id = "1234567890" + project_id = "test-project" + location = "us-central1" + + auth_header, url = vertex_base._check_custom_proxy( + api_base=psc_api_base, + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="predict", + stream=False, + auth_header="test-token", + url="", + model=endpoint_id, + vertex_project=project_id, + vertex_location=location, + vertex_api_version="v1", + ) + + # rstrip('/') should remove the trailing slash + expected_url = f"{psc_api_base.rstrip('/')}/v1/projects/{project_id}/locations/{location}/endpoints/{endpoint_id}:predict" + assert ( + url == expected_url + ), f"Expected {expected_url}, but got {url}" + + def test_standard_proxy_with_googleapis(self): + """Test that standard proxies with googleapis.com in URL use simple format""" + vertex_base = VertexBase() + proxy_api_base = "https://my-proxy.googleapis.com" + endpoint_id = "gemini-pro" # Not numeric + project_id = "test-project" + location = "us-central1" + + auth_header, url = vertex_base._check_custom_proxy( + api_base=proxy_api_base, + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="generateContent", + stream=False, + auth_header="test-token", + url="", + model=endpoint_id, + vertex_project=project_id, + vertex_location=location, + vertex_api_version="v1", + ) + + # Should use simple format: api_base:endpoint + expected_url = f"{proxy_api_base}:generateContent" + assert ( + url == expected_url + ), f"Expected {expected_url}, but got {url}" + + def test_custom_proxy_with_numeric_model(self): + """Test that numeric model IDs trigger PSC-style URL construction""" + vertex_base = VertexBase() + proxy_api_base = "https://my-custom-proxy.example.com" + endpoint_id = "9876543210" # Numeric endpoint ID + project_id = "test-project" + location = "us-central1" + + auth_header, url = vertex_base._check_custom_proxy( + api_base=proxy_api_base, + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="predict", + stream=False, + auth_header="test-token", + url="", + model=endpoint_id, + vertex_project=project_id, + vertex_location=location, + vertex_api_version="v1", + ) + + # Numeric model should trigger full path construction + expected_url = f"{proxy_api_base}/v1/projects/{project_id}/locations/{location}/endpoints/{endpoint_id}:predict" + assert ( + url == expected_url + ), f"Expected {expected_url}, but got {url}" + + def test_no_api_base_returns_original_url(self): + """Test that when api_base is None, the original URL is returned""" + vertex_base = VertexBase() + original_url = "https://us-central1-aiplatform.googleapis.com/v1/projects/test/locations/us-central1/publishers/google/models/gemini-pro:generateContent" + + auth_header, url = vertex_base._check_custom_proxy( + api_base=None, + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="generateContent", + stream=False, + auth_header="test-token", + url=original_url, + model="gemini-pro", + vertex_project="test-project", + vertex_location="us-central1", + vertex_api_version="v1", + ) + + # When api_base is None, original URL should be returned unchanged + assert url == original_url, f"Expected {original_url}, but got {url}" + + def test_auth_header_preserved(self): + """Test that auth_header is properly preserved""" + vertex_base = VertexBase() + psc_api_base = "http://10.96.32.8" + test_auth_header = "Bearer test-token-12345" + + auth_header, url = vertex_base._check_custom_proxy( + api_base=psc_api_base, + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="predict", + stream=False, + auth_header=test_auth_header, + url="", + model="1234567890", + vertex_project="test-project", + vertex_location="us-central1", + vertex_api_version="v1", + ) + + assert ( + auth_header == test_auth_header + ), f"Auth header should be preserved, got {auth_header}" + From 62f2bb0ed0e8e4cd8f4633e2e6ea742d03a2d15a Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:10:35 -0700 Subject: [PATCH 02/23] add TextEmbeddingBGEInput --- litellm/llms/vertex_ai/vertex_embeddings/types.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/litellm/llms/vertex_ai/vertex_embeddings/types.py b/litellm/llms/vertex_ai/vertex_embeddings/types.py index 7f85ea46f31..fa9794d79a5 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/types.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/types.py @@ -25,6 +25,12 @@ class TextEmbeddingInput(TypedDict, total=False): title: Optional[str] +class TextEmbeddingBGEInput(TypedDict, total=False): + prompt: str + task_type: Optional[TaskType] + title: Optional[str] + + # Fine-tuned models require a different input format # Ref: https://console.cloud.google.com/vertex-ai/model-garden?hl=en&project=adroit-crow-413218&pageState=(%22galleryStateKey%22:(%22f%22:(%22g%22:%5B%5D,%22o%22:%5B%5D),%22s%22:%22%22)) class TextEmbeddingFineTunedInput(TypedDict, total=False): @@ -44,7 +50,7 @@ class EmbeddingParameters(TypedDict, total=False): class VertexEmbeddingRequest(TypedDict, total=False): - instances: Union[List[TextEmbeddingInput], List[TextEmbeddingFineTunedInput]] + instances: Union[List[TextEmbeddingInput], List[TextEmbeddingBGEInput], List[TextEmbeddingFineTunedInput]] parameters: Optional[Union[EmbeddingParameters, TextEmbeddingFineTunedParameters]] From 39e750d3b2c3d68a88ee8de3abbde0aef5ad653f Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:10:45 -0700 Subject: [PATCH 03/23] add VertexBGEConfig --- .../llms/vertex_ai/vertex_embeddings/bge.py | 100 ++++++++++++++++++ 1 file changed, 100 insertions(+) create mode 100644 litellm/llms/vertex_ai/vertex_embeddings/bge.py diff --git a/litellm/llms/vertex_ai/vertex_embeddings/bge.py b/litellm/llms/vertex_ai/vertex_embeddings/bge.py new file mode 100644 index 00000000000..401f7ebd907 --- /dev/null +++ b/litellm/llms/vertex_ai/vertex_embeddings/bge.py @@ -0,0 +1,100 @@ +""" +Vertex AI BGE (BAAI General Embedding) Configuration + +BGE models deployed on Vertex AI require different input format: +- Use "prompt" instead of "content" as the input field +""" + +from typing import List, Optional, Union + +from .types import ( + EmbeddingParameters, + TaskType, + TextEmbeddingBGEInput, + VertexEmbeddingRequest, +) + + +class VertexBGEConfig: + """ + Configuration and transformation logic for BGE models on Vertex AI. + + BGE (BAAI General Embedding) models use a different request format + where the input field is named "prompt" instead of "content". + """ + + @staticmethod + def is_bge_model(model: str) -> bool: + """ + Check if the model is a BGE (BAAI General Embedding) model. + + Args: + model: The model name + + Returns: + bool: True if the model is a BGE model + """ + return "bge" in model.lower() + + @staticmethod + def transform_request( + input: Union[list, str], optional_params: dict, model: str + ) -> VertexEmbeddingRequest: + """ + Transforms an OpenAI request to a Vertex BGE embedding request. + + BGE models use "prompt" instead of "content" as the input field. + + Args: + input: The input text(s) to embed + optional_params: Optional parameters for the request + model: The model name + + Returns: + VertexEmbeddingRequest: The transformed request + """ + vertex_request: VertexEmbeddingRequest = VertexEmbeddingRequest() + vertex_text_embedding_input_list: List[TextEmbeddingBGEInput] = [] + task_type: Optional[TaskType] = optional_params.get("task_type") + title = optional_params.get("title") + + if isinstance(input, str): + input = [input] + + for text in input: + embedding_input = VertexBGEConfig._create_embedding_input( + prompt=text, task_type=task_type, title=title + ) + vertex_text_embedding_input_list.append(embedding_input) + + vertex_request["instances"] = vertex_text_embedding_input_list + vertex_request["parameters"] = EmbeddingParameters(**optional_params) + + return vertex_request + + @staticmethod + def _create_embedding_input( + prompt: str, + task_type: Optional[TaskType] = None, + title: Optional[str] = None, + ) -> TextEmbeddingBGEInput: + """ + Creates a TextEmbeddingBGEInput object for BGE models. + + BGE models use "prompt" instead of "content" as the input field. + + Args: + prompt: The prompt to be embedded + task_type: The type of task to be performed + title: The title of the document to be embedded + + Returns: + TextEmbeddingBGEInput: A TextEmbeddingBGEInput object + """ + text_embedding_input = TextEmbeddingBGEInput(prompt=prompt) + if task_type is not None: + text_embedding_input["task_type"] = task_type + if title is not None: + text_embedding_input["title"] = title + return text_embedding_input + From 3293ac8a3d282717b32e1b4aa562caede70151eb Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:11:01 -0700 Subject: [PATCH 04/23] add BGE handling --- .../llms/vertex_ai/vertex_embeddings/transformation.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index caaf00e199e..7bbe13e3597 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -5,6 +5,7 @@ from pydantic import BaseModel from litellm.types.utils import EmbeddingResponse, Usage +from .bge import VertexBGEConfig from .types import * @@ -109,6 +110,11 @@ class VertexAITextEmbeddingConfig(BaseModel): return self._transform_openai_request_to_fine_tuned_embedding_request( input, optional_params, model ) + + if VertexBGEConfig.is_bge_model(model): + return VertexBGEConfig.transform_request( + input=input, optional_params=optional_params, model=model + ) vertex_request: VertexEmbeddingRequest = VertexEmbeddingRequest() vertex_text_embedding_input_list: List[TextEmbeddingInput] = [] @@ -186,8 +192,8 @@ class VertexAITextEmbeddingConfig(BaseModel): Args: content (str): The content to be embedded. - task_type (Optional[TaskType]): The type of task to be performed". - title (Optional[str]): The title of the document to be embedded + task_type (Optional[TaskType]): The type of task to be performed. + title (Optional[str]): The title of the document to be embedded. Returns: TextEmbeddingInput: A TextEmbeddingInput object. From f2befcf6572c58d97ef977d496afadc9d3766deb Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:12:29 -0700 Subject: [PATCH 05/23] test_vertex_ai_bge_embedding_with_custom_api_base --- .../llms/vertex_ai/test_bge_embedding.py | 106 ++++++++++++++++++ 1 file changed, 106 insertions(+) create mode 100644 tests/test_litellm/llms/vertex_ai/test_bge_embedding.py diff --git a/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py b/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py new file mode 100644 index 00000000000..d7bd07f53ef --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py @@ -0,0 +1,106 @@ +""" +Test BGE embeddings with Vertex AI using custom api_base. + +This test ensures that BGE embeddings work correctly with Vertex AI +and that the request body is properly formatted. +""" + +import json +import os +import sys +from unittest.mock import MagicMock, patch + +sys.path.insert( + 0, os.path.abspath("../../../..") +) + +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler + + +def test_vertex_ai_bge_embedding_with_custom_api_base(): + """ + Test Vertex AI BGE embeddings with custom api_base. + + This test verifies that when using a BGE model with Vertex AI and + a custom api_base, the request is properly formatted and sent to + the correct endpoint. + """ + client = HTTPHandler() + + def mock_auth_token(*args, **kwargs): + return "fake-token", "fake-project" + + with patch.object(client, "post") as mock_post, patch( + "litellm.llms.vertex_ai.vertex_embeddings.embedding_handler.VertexEmbedding._ensure_access_token", + side_effect=mock_auth_token + ): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "predictions": [ + { + "embeddings": { + "values": [0.1, 0.2, 0.3, 0.4, 0.5], + "statistics": {"token_count": 2} + } + }, + { + "embeddings": { + "values": [0.6, 0.7, 0.8, 0.9, 1.0], + "statistics": {"token_count": 2} + } + } + ] + } + mock_post.return_value = mock_response + + response = litellm.embedding( + model="vertex_ai/bge-small-en-v1.5", + input=["Hello", "World"], + api_base="http://10.96.32.8", + client=client + ) + + mock_post.assert_called_once() + + call_args = mock_post.call_args + kwargs = call_args.kwargs if hasattr(call_args, 'kwargs') else call_args[1] + + if "url" in kwargs: + api_url_called = kwargs["url"] + elif len(call_args[0]) > 0: + api_url_called = call_args[0][0] + else: + api_url_called = "Unknown" + + # Vertex AI may use 'json' or 'data' parameter + if "json" in kwargs: + request_data = kwargs["json"] + elif "data" in kwargs: + request_data = json.loads(kwargs["data"]) + else: + request_data = {} + + print("\n" + "="*50) + print("Mock Request Body Received:") + print("="*50) + print(json.dumps(request_data, indent=2)) + print("="*50) + print(f"API Base: {api_url_called}") + print("="*50 + "\n") + + assert "instances" in request_data + assert len(request_data["instances"]) == 2 + # BGE models should use "prompt" instead of "content" + assert "prompt" in request_data["instances"][0] + assert request_data["instances"][0]["prompt"] == "Hello" + assert "prompt" in request_data["instances"][1] + assert request_data["instances"][1]["prompt"] == "World" + + assert isinstance(response.data, list) + assert len(response.data) == 2 + assert "embedding" in response.data[0] + From c821acd61a1eef41589df3410b8d2d23d9395ab0 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:14:25 -0700 Subject: [PATCH 06/23] fix request transform vertex BGE --- .../llms/vertex_ai/vertex_embeddings/bge.py | 54 ++++++++++++++++++- .../vertex_embeddings/transformation.py | 5 ++ 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_embeddings/bge.py b/litellm/llms/vertex_ai/vertex_embeddings/bge.py index 401f7ebd907..1bfa362ee98 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/bge.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/bge.py @@ -1,12 +1,15 @@ """ Vertex AI BGE (BAAI General Embedding) Configuration -BGE models deployed on Vertex AI require different input format: -- Use "prompt" instead of "content" as the input field +BGE models deployed on Vertex AI require different input/output format: +- Request: Use "prompt" instead of "content" as the input field +- Response: Embeddings are returned directly as arrays, not wrapped in objects """ from typing import List, Optional, Union +from litellm.types.utils import EmbeddingResponse, Usage + from .types import ( EmbeddingParameters, TaskType, @@ -98,3 +101,50 @@ class VertexBGEConfig: text_embedding_input["title"] = title return text_embedding_input + @staticmethod + def transform_response( + response: dict, model: str, model_response: EmbeddingResponse + ) -> EmbeddingResponse: + """ + Transforms a Vertex BGE embedding response to OpenAI format. + + BGE models return embeddings directly as arrays in predictions: + { + "predictions": [ + [0.002, 0.021, ...], + [0.003, 0.022, ...] + ] + } + + Args: + response: The raw response from Vertex AI + model: The model name + model_response: The EmbeddingResponse object to populate + + Returns: + EmbeddingResponse: The transformed response in OpenAI format + """ + _predictions = response["predictions"] + + embedding_response = [] + # BGE models don't return token counts, so we estimate or set to 0 + input_tokens = 0 + + for idx, embedding_values in enumerate(_predictions): + embedding_response.append( + { + "object": "embedding", + "index": idx, + "embedding": embedding_values, + } + ) + + model_response.object = "list" + model_response.data = embedding_response + model_response.model = model + usage = Usage( + prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens + ) + setattr(model_response, "usage", usage) + return model_response + diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index 7bbe13e3597..77da3ce7c01 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -215,6 +215,11 @@ class VertexAITextEmbeddingConfig(BaseModel): return self._transform_vertex_response_to_openai_for_fine_tuned_models( response, model, model_response ) + + if VertexBGEConfig.is_bge_model(model): + return VertexBGEConfig.transform_response( + response=response, model=model, model_response=model_response + ) _predictions = response["predictions"] From 8957770e68489ea0e36fda49cf0f4114753433f4 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:14:34 -0700 Subject: [PATCH 07/23] test_vertex_ai_bge_embedding_with_custom_api_base --- .../llms/vertex_ai/test_bge_embedding.py | 21 +++++++------------ 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py b/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py index d7bd07f53ef..636df93b026 100644 --- a/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py +++ b/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py @@ -39,21 +39,16 @@ def test_vertex_ai_bge_embedding_with_custom_api_base(): ): mock_response = MagicMock() mock_response.status_code = 200 + # BGE models return embeddings directly as arrays, not wrapped in objects mock_response.json.return_value = { "predictions": [ - { - "embeddings": { - "values": [0.1, 0.2, 0.3, 0.4, 0.5], - "statistics": {"token_count": 2} - } - }, - { - "embeddings": { - "values": [0.6, 0.7, 0.8, 0.9, 1.0], - "statistics": {"token_count": 2} - } - } - ] + [0.1, 0.2, 0.3, 0.4, 0.5], + [0.6, 0.7, 0.8, 0.9, 1.0] + ], + "deployedModelId": "849506872875548672", + "model": "projects/1060139831167/locations/us-central1/models/baai_bge-small-en-v1.5", + "modelDisplayName": "baai_bge-small-en-v1.5", + "modelVersionId": "1" } mock_post.return_value = mock_response From 0abf450b7e0f768b3295a815fbac25c42b8cca64 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:15:58 -0700 Subject: [PATCH 08/23] tes BGE --- .../llms/vertex_ai/vertex_embeddings/bge.py | 15 +++ .../test_bge_response_transformation.py | 93 +++++++++++++++++++ 2 files changed, 108 insertions(+) create mode 100644 tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py diff --git a/litellm/llms/vertex_ai/vertex_embeddings/bge.py b/litellm/llms/vertex_ai/vertex_embeddings/bge.py index 1bfa362ee98..b8979f55880 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/bge.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/bge.py @@ -123,14 +123,29 @@ class VertexBGEConfig: Returns: EmbeddingResponse: The transformed response in OpenAI format + + Raises: + KeyError: If response doesn't contain 'predictions' + ValueError: If predictions is not a list or contains invalid data """ + if "predictions" not in response: + raise KeyError("Response missing 'predictions' field") + _predictions = response["predictions"] + + if not isinstance(_predictions, list): + raise ValueError(f"Expected 'predictions' to be a list, got {type(_predictions)}") embedding_response = [] # BGE models don't return token counts, so we estimate or set to 0 input_tokens = 0 for idx, embedding_values in enumerate(_predictions): + if not isinstance(embedding_values, list): + raise ValueError( + f"Expected embedding at index {idx} to be a list, got {type(embedding_values)}" + ) + embedding_response.append( { "object": "embedding", diff --git a/tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py b/tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py new file mode 100644 index 00000000000..a3f28678229 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py @@ -0,0 +1,93 @@ +""" +Test BGE response transformation validation. + +This test verifies that the BGE response transformer properly validates +and handles different response formats. +""" + +import os +import sys + +sys.path.insert( + 0, os.path.abspath("../../../..") +) + +import pytest + +from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig +from litellm.types.utils import EmbeddingResponse + + +def test_bge_response_transformation_success(): + """ + Test successful BGE response transformation. + + Verifies that a valid BGE response is properly transformed + to OpenAI format. + """ + response = { + "predictions": [ + [0.1, 0.2, 0.3], + [0.4, 0.5, 0.6] + ], + "deployedModelId": "123456", + "model": "projects/test/models/bge-base" + } + + model_response = EmbeddingResponse() + result = VertexBGEConfig.transform_response( + response=response, + model="bge-small-en-v1.5", + model_response=model_response + ) + + assert result.object == "list" + assert len(result.data) == 2 + assert result.data[0]["embedding"] == [0.1, 0.2, 0.3] + assert result.data[1]["embedding"] == [0.4, 0.5, 0.6] + assert result.data[0]["index"] == 0 + assert result.data[1]["index"] == 1 + assert result.model == "bge-small-en-v1.5" + + +def test_bge_response_missing_predictions(): + """ + Test BGE response transformation with missing predictions field. + + Verifies that a KeyError is raised when the response doesn't + contain the required 'predictions' field. + """ + response = { + "deployedModelId": "123456", + "model": "projects/test/models/bge-base" + } + + model_response = EmbeddingResponse() + + with pytest.raises(KeyError, match="Response missing 'predictions' field"): + VertexBGEConfig.transform_response( + response=response, + model="bge-small-en-v1.5", + model_response=model_response + ) + + +def test_bge_response_invalid_predictions_type(): + """ + Test BGE response transformation with invalid predictions type. + + Verifies that a ValueError is raised when predictions is not a list. + """ + response = { + "predictions": "not-a-list" + } + + model_response = EmbeddingResponse() + + with pytest.raises(ValueError, match="Expected 'predictions' to be a list"): + VertexBGEConfig.transform_response( + response=response, + model="bge-small-en-v1.5", + model_response=model_response + ) + From 58d9531869f9588ca7f473f2edca60b170a65f4a Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:38:48 -0700 Subject: [PATCH 09/23] test_is_bge_model_detection --- .../test_bge_response_transformation.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py b/tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py index a3f28678229..20150501adf 100644 --- a/tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py @@ -18,6 +18,24 @@ from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig from litellm.types.utils import EmbeddingResponse +def test_is_bge_model_detection(): + """ + Test BGE model detection for post-provider-split patterns. + + After main.py splits the provider, model strings are passed without the provider prefix. + Model name transformation (bge/ -> numeric ID) is handled in common_utils._get_vertex_url(). + """ + # Should detect BGE models (after provider split) + assert VertexBGEConfig.is_bge_model("bge-small-en-v1.5") is True + assert VertexBGEConfig.is_bge_model("bge/204379420394258432") is True + assert VertexBGEConfig.is_bge_model("BGE-large-en-v1.5") is True # case insensitive + + # Should not detect non-BGE models + assert VertexBGEConfig.is_bge_model("textembedding-gecko") is False + assert VertexBGEConfig.is_bge_model("gemma") is False + assert VertexBGEConfig.is_bge_model("123456789") is False + + def test_bge_response_transformation_success(): """ Test successful BGE response transformation. From 88b2cfc789665a0ea174a383ae10576fe7225725 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:41:33 -0700 Subject: [PATCH 10/23] docs cleanup --- docs/my-website/docs/providers/vertex.md | 509 ----------------- .../docs/providers/vertex_embedding.md | 511 ++++++++++++++++++ docs/my-website/sidebars.js | 1 + 3 files changed, 512 insertions(+), 509 deletions(-) create mode 100644 docs/my-website/docs/providers/vertex_embedding.md diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 8e333b69ef7..5df63582446 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -2089,515 +2089,6 @@ curl http://0.0.0.0:4000/v1/chat/completions \ | code-gecko@latest| `completion('code-gecko@latest', messages)` | -## **Embedding Models** - -#### Usage - Embedding - - - - -```python -import litellm -from litellm import embedding -litellm.vertex_project = "hardy-device-38811" # Your Project ID -litellm.vertex_location = "us-central1" # proj location - -response = embedding( - model="vertex_ai/textembedding-gecko", - input=["good morning from litellm"], -) -print(response) -``` - - - - - -1. Add model to config.yaml -```yaml -model_list: - - model_name: snowflake-arctic-embed-m-long-1731622468876 - litellm_params: - model: vertex_ai/ - vertex_project: "adroit-crow-413218" - vertex_location: "us-central1" - vertex_credentials: adroit-crow-413218-a956eef1a2a8.json - -litellm_settings: - drop_params: True -``` - -2. Start Proxy - -``` -$ litellm --config /path/to/config.yaml -``` - -3. Make Request using OpenAI Python SDK, Langchain Python SDK - -```python -import openai - -client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") - -response = client.embeddings.create( - model="snowflake-arctic-embed-m-long-1731622468876", - input = ["good morning from litellm", "this is another item"], -) - -print(response) -``` - - - - - -#### Supported Embedding Models -All models listed [here](https://github.com/BerriAI/litellm/blob/57f37f743886a0249f630a6792d49dffc2c5d9b7/model_prices_and_context_window.json#L835) are supported - -| Model Name | Function Call | -|--------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| text-embedding-004 | `embedding(model="vertex_ai/text-embedding-004", input)` | -| text-multilingual-embedding-002 | `embedding(model="vertex_ai/text-multilingual-embedding-002", input)` | -| textembedding-gecko | `embedding(model="vertex_ai/textembedding-gecko", input)` | -| textembedding-gecko-multilingual | `embedding(model="vertex_ai/textembedding-gecko-multilingual", input)` | -| textembedding-gecko-multilingual@001 | `embedding(model="vertex_ai/textembedding-gecko-multilingual@001", input)` | -| textembedding-gecko@001 | `embedding(model="vertex_ai/textembedding-gecko@001", input)` | -| textembedding-gecko@003 | `embedding(model="vertex_ai/textembedding-gecko@003", input)` | -| text-embedding-preview-0409 | `embedding(model="vertex_ai/text-embedding-preview-0409", input)` | -| text-multilingual-embedding-preview-0409 | `embedding(model="vertex_ai/text-multilingual-embedding-preview-0409", input)` | -| Fine-tuned OR Custom Embedding models | `embedding(model="vertex_ai/", input)` | - -### Supported OpenAI (Unified) Params - -| [param](../embedding/supported_embedding.md#input-params-for-litellmembedding) | type | [vertex equivalent](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api) | -|-------|-------------|--------------------| -| `input` | **string or List[string]** | `instances` | -| `dimensions` | **int** | `output_dimensionality` | -| `input_type` | **Literal["RETRIEVAL_QUERY","RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING", "FACT_VERIFICATION"]** | `task_type` | - -#### Usage with OpenAI (Unified) Params - - - - - -```python -response = litellm.embedding( - model="vertex_ai/text-embedding-004", - input=["good morning from litellm", "gm"] - input_type = "RETRIEVAL_DOCUMENT", - dimensions=1, -) -``` - - - - -```python -import openai - -client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") - -response = client.embeddings.create( - model="text-embedding-004", - input = ["good morning from litellm", "gm"], - dimensions=1, - extra_body = { - "input_type": "RETRIEVAL_QUERY", - } -) - -print(response) -``` - - - - -### Supported Vertex Specific Params - -| param | type | -|-------|-------------| -| `auto_truncate` | **bool** | -| `task_type` | **Literal["RETRIEVAL_QUERY","RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING", "FACT_VERIFICATION"]** | -| `title` | **str** | - -#### Usage with Vertex Specific Params (Use `task_type` and `title`) - -You can pass any vertex specific params to the embedding model. Just pass them to the embedding function like this: - -[Relevant Vertex AI doc with all embedding params](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api#request_body) - - - - -```python -response = litellm.embedding( - model="vertex_ai/text-embedding-004", - input=["good morning from litellm", "gm"] - task_type = "RETRIEVAL_DOCUMENT", - title = "test", - dimensions=1, - auto_truncate=True, -) -``` - - - - -```python -import openai - -client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") - -response = client.embeddings.create( - model="text-embedding-004", - input = ["good morning from litellm", "gm"], - dimensions=1, - extra_body = { - "task_type": "RETRIEVAL_QUERY", - "auto_truncate": True, - "title": "test", - } -) - -print(response) -``` - - - -## **Multi-Modal Embeddings** - - -Known Limitations: -- Only supports 1 image / video / image per request -- Only supports GCS or base64 encoded images / videos - -### Usage - - - - -Using GCS Images - -```python -response = await litellm.aembedding( - model="vertex_ai/multimodalembedding@001", - input="gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png" # will be sent as a gcs image -) -``` - -Using base 64 encoded images - -```python -response = await litellm.aembedding( - model="vertex_ai/multimodalembedding@001", - input="data:image/jpeg;base64,..." # will be sent as a base64 encoded image -) -``` - - - - -1. Add model to config.yaml -```yaml -model_list: - - model_name: multimodalembedding@001 - litellm_params: - model: vertex_ai/multimodalembedding@001 - vertex_project: "adroit-crow-413218" - vertex_location: "us-central1" - vertex_credentials: adroit-crow-413218-a956eef1a2a8.json - -litellm_settings: - drop_params: True -``` - -2. Start Proxy - -``` -$ litellm --config /path/to/config.yaml -``` - -3. Make Request use OpenAI Python SDK, Langchain Python SDK - - - - - - -Requests with GCS Image / Video URI - -```python -import openai - -client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") - -# # request sent to model set on litellm proxy, `litellm --model` -response = client.embeddings.create( - model="multimodalembedding@001", - input = "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png", -) - -print(response) -``` - -Requests with base64 encoded images - -```python -import openai - -client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") - -# # request sent to model set on litellm proxy, `litellm --model` -response = client.embeddings.create( - model="multimodalembedding@001", - input = "data:image/jpeg;base64,...", -) - -print(response) -``` - - - - - -Requests with GCS Image / Video URI -```python -from langchain_openai import OpenAIEmbeddings - -embeddings_models = "multimodalembedding@001" - -embeddings = OpenAIEmbeddings( - model="multimodalembedding@001", - base_url="http://0.0.0.0:4000", - api_key="sk-1234", # type: ignore -) - - -query_result = embeddings.embed_query( - "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png" -) -print(query_result) - -``` - -Requests with base64 encoded images - -```python -from langchain_openai import OpenAIEmbeddings - -embeddings_models = "multimodalembedding@001" - -embeddings = OpenAIEmbeddings( - model="multimodalembedding@001", - base_url="http://0.0.0.0:4000", - api_key="sk-1234", # type: ignore -) - - -query_result = embeddings.embed_query( - "data:image/jpeg;base64,..." -) -print(query_result) - -``` - - - - - - - - - -1. Add model to config.yaml -```yaml -default_vertex_config: - vertex_project: "adroit-crow-413218" - vertex_location: "us-central1" - vertex_credentials: adroit-crow-413218-a956eef1a2a8.json -``` - -2. Start Proxy - -``` -$ litellm --config /path/to/config.yaml -``` - -3. Make Request use OpenAI Python SDK - -```python -import vertexai - -from vertexai.vision_models import Image, MultiModalEmbeddingModel, Video -from vertexai.vision_models import VideoSegmentConfig -from google.auth.credentials import Credentials - - -LITELLM_PROXY_API_KEY = "sk-1234" -LITELLM_PROXY_BASE = "http://0.0.0.0:4000/vertex-ai" - -import datetime - -class CredentialsWrapper(Credentials): - def __init__(self, token=None): - super().__init__() - self.token = token - self.expiry = None # or set to a future date if needed - - def refresh(self, request): - pass - - def apply(self, headers, token=None): - headers['Authorization'] = f'Bearer {self.token}' - - @property - def expired(self): - return False # Always consider the token as non-expired - - @property - def valid(self): - return True # Always consider the credentials as valid - -credentials = CredentialsWrapper(token=LITELLM_PROXY_API_KEY) - -vertexai.init( - project="adroit-crow-413218", - location="us-central1", - api_endpoint=LITELLM_PROXY_BASE, - credentials = credentials, - api_transport="rest", - -) - -model = MultiModalEmbeddingModel.from_pretrained("multimodalembedding") -image = Image.load_from_file( - "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png" -) - -embeddings = model.get_embeddings( - image=image, - contextual_text="Colosseum", - dimension=1408, -) -print(f"Image Embedding: {embeddings.image_embedding}") -print(f"Text Embedding: {embeddings.text_embedding}") -``` - - - - - -### Text + Image + Video Embeddings - - - - -Text + Image - -```python -response = await litellm.aembedding( - model="vertex_ai/multimodalembedding@001", - input=["hey", "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"] # will be sent as a gcs image -) -``` - -Text + Video - -```python -response = await litellm.aembedding( - model="vertex_ai/multimodalembedding@001", - input=["hey", "gs://my-bucket/embeddings/supermarket-video.mp4"] # will be sent as a gcs image -) -``` - -Image + Video - -```python -response = await litellm.aembedding( - model="vertex_ai/multimodalembedding@001", - input=["gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png", "gs://my-bucket/embeddings/supermarket-video.mp4"] # will be sent as a gcs image -) -``` - - - - - -1. Add model to config.yaml -```yaml -model_list: - - model_name: multimodalembedding@001 - litellm_params: - model: vertex_ai/multimodalembedding@001 - vertex_project: "adroit-crow-413218" - vertex_location: "us-central1" - vertex_credentials: adroit-crow-413218-a956eef1a2a8.json - -litellm_settings: - drop_params: True -``` - -2. Start Proxy - -``` -$ litellm --config /path/to/config.yaml -``` - -3. Make Request use OpenAI Python SDK, Langchain Python SDK - - -Text + Image - -```python -import openai - -client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") - -# # request sent to model set on litellm proxy, `litellm --model` -response = client.embeddings.create( - model="multimodalembedding@001", - input = ["hey", "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"], -) - -print(response) -``` - -Text + Video -```python -import openai - -client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") - -# # request sent to model set on litellm proxy, `litellm --model` -response = client.embeddings.create( - model="multimodalembedding@001", - input = ["hey", "gs://my-bucket/embeddings/supermarket-video.mp4"], -) - -print(response) -``` - -Image + Video -```python -import openai - -client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") - -# # request sent to model set on litellm proxy, `litellm --model` -response = client.embeddings.create( - model="multimodalembedding@001", - input = ["gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png", "gs://my-bucket/embeddings/supermarket-video.mp4"], -) - -print(response) -``` - - - - - ## **Gemini TTS (Text-to-Speech) Audio Output** :::info diff --git a/docs/my-website/docs/providers/vertex_embedding.md b/docs/my-website/docs/providers/vertex_embedding.md new file mode 100644 index 00000000000..25580935387 --- /dev/null +++ b/docs/my-website/docs/providers/vertex_embedding.md @@ -0,0 +1,511 @@ +import Image from '@theme/IdealImage'; +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Vertex AI Embedding + +## Usage - Embedding + + + + +```python +import litellm +from litellm import embedding +litellm.vertex_project = "hardy-device-38811" # Your Project ID +litellm.vertex_location = "us-central1" # proj location + +response = embedding( + model="vertex_ai/textembedding-gecko", + input=["good morning from litellm"], +) +print(response) +``` + + + + + +1. Add model to config.yaml +```yaml +model_list: + - model_name: snowflake-arctic-embed-m-long-1731622468876 + litellm_params: + model: vertex_ai/ + vertex_project: "adroit-crow-413218" + vertex_location: "us-central1" + vertex_credentials: adroit-crow-413218-a956eef1a2a8.json + +litellm_settings: + drop_params: True +``` + +2. Start Proxy + +``` +$ litellm --config /path/to/config.yaml +``` + +3. Make Request using OpenAI Python SDK, Langchain Python SDK + +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +response = client.embeddings.create( + model="snowflake-arctic-embed-m-long-1731622468876", + input = ["good morning from litellm", "this is another item"], +) + +print(response) +``` + + + + + +#### Supported Embedding Models +All models listed [here](https://github.com/BerriAI/litellm/blob/57f37f743886a0249f630a6792d49dffc2c5d9b7/model_prices_and_context_window.json#L835) are supported + +| Model Name | Function Call | +|--------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| text-embedding-004 | `embedding(model="vertex_ai/text-embedding-004", input)` | +| text-multilingual-embedding-002 | `embedding(model="vertex_ai/text-multilingual-embedding-002", input)` | +| textembedding-gecko | `embedding(model="vertex_ai/textembedding-gecko", input)` | +| textembedding-gecko-multilingual | `embedding(model="vertex_ai/textembedding-gecko-multilingual", input)` | +| textembedding-gecko-multilingual@001 | `embedding(model="vertex_ai/textembedding-gecko-multilingual@001", input)` | +| textembedding-gecko@001 | `embedding(model="vertex_ai/textembedding-gecko@001", input)` | +| textembedding-gecko@003 | `embedding(model="vertex_ai/textembedding-gecko@003", input)` | +| text-embedding-preview-0409 | `embedding(model="vertex_ai/text-embedding-preview-0409", input)` | +| text-multilingual-embedding-preview-0409 | `embedding(model="vertex_ai/text-multilingual-embedding-preview-0409", input)` | +| Fine-tuned OR Custom Embedding models | `embedding(model="vertex_ai/", input)` | + +### Supported OpenAI (Unified) Params + +| [param](../embedding/supported_embedding.md#input-params-for-litellmembedding) | type | [vertex equivalent](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api) | +|-------|-------------|--------------------| +| `input` | **string or List[string]** | `instances` | +| `dimensions` | **int** | `output_dimensionality` | +| `input_type` | **Literal["RETRIEVAL_QUERY","RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING", "FACT_VERIFICATION"]** | `task_type` | + +#### Usage with OpenAI (Unified) Params + + + + + +```python +response = litellm.embedding( + model="vertex_ai/text-embedding-004", + input=["good morning from litellm", "gm"] + input_type = "RETRIEVAL_DOCUMENT", + dimensions=1, +) +``` + + + + +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +response = client.embeddings.create( + model="text-embedding-004", + input = ["good morning from litellm", "gm"], + dimensions=1, + extra_body = { + "input_type": "RETRIEVAL_QUERY", + } +) + +print(response) +``` + + + + +### Supported Vertex Specific Params + +| param | type | +|-------|-------------| +| `auto_truncate` | **bool** | +| `task_type` | **Literal["RETRIEVAL_QUERY","RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING", "FACT_VERIFICATION"]** | +| `title` | **str** | + +#### Usage with Vertex Specific Params (Use `task_type` and `title`) + +You can pass any vertex specific params to the embedding model. Just pass them to the embedding function like this: + +[Relevant Vertex AI doc with all embedding params](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api#request_body) + + + + +```python +response = litellm.embedding( + model="vertex_ai/text-embedding-004", + input=["good morning from litellm", "gm"] + task_type = "RETRIEVAL_DOCUMENT", + title = "test", + dimensions=1, + auto_truncate=True, +) +``` + + + + +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +response = client.embeddings.create( + model="text-embedding-004", + input = ["good morning from litellm", "gm"], + dimensions=1, + extra_body = { + "task_type": "RETRIEVAL_QUERY", + "auto_truncate": True, + "title": "test", + } +) + +print(response) +``` + + + +## **Multi-Modal Embeddings** + + +Known Limitations: +- Only supports 1 image / video / image per request +- Only supports GCS or base64 encoded images / videos + +### Usage + + + + +Using GCS Images + +```python +response = await litellm.aembedding( + model="vertex_ai/multimodalembedding@001", + input="gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png" # will be sent as a gcs image +) +``` + +Using base 64 encoded images + +```python +response = await litellm.aembedding( + model="vertex_ai/multimodalembedding@001", + input="data:image/jpeg;base64,..." # will be sent as a base64 encoded image +) +``` + + + + +1. Add model to config.yaml +```yaml +model_list: + - model_name: multimodalembedding@001 + litellm_params: + model: vertex_ai/multimodalembedding@001 + vertex_project: "adroit-crow-413218" + vertex_location: "us-central1" + vertex_credentials: adroit-crow-413218-a956eef1a2a8.json + +litellm_settings: + drop_params: True +``` + +2. Start Proxy + +``` +$ litellm --config /path/to/config.yaml +``` + +3. Make Request use OpenAI Python SDK, Langchain Python SDK + + + + + + +Requests with GCS Image / Video URI + +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +# # request sent to model set on litellm proxy, `litellm --model` +response = client.embeddings.create( + model="multimodalembedding@001", + input = "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png", +) + +print(response) +``` + +Requests with base64 encoded images + +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +# # request sent to model set on litellm proxy, `litellm --model` +response = client.embeddings.create( + model="multimodalembedding@001", + input = "data:image/jpeg;base64,...", +) + +print(response) +``` + + + + + +Requests with GCS Image / Video URI +```python +from langchain_openai import OpenAIEmbeddings + +embeddings_models = "multimodalembedding@001" + +embeddings = OpenAIEmbeddings( + model="multimodalembedding@001", + base_url="http://0.0.0.0:4000", + api_key="sk-1234", # type: ignore +) + + +query_result = embeddings.embed_query( + "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png" +) +print(query_result) + +``` + +Requests with base64 encoded images + +```python +from langchain_openai import OpenAIEmbeddings + +embeddings_models = "multimodalembedding@001" + +embeddings = OpenAIEmbeddings( + model="multimodalembedding@001", + base_url="http://0.0.0.0:4000", + api_key="sk-1234", # type: ignore +) + + +query_result = embeddings.embed_query( + "data:image/jpeg;base64,..." +) +print(query_result) + +``` + + + + + + + + + +1. Add model to config.yaml +```yaml +default_vertex_config: + vertex_project: "adroit-crow-413218" + vertex_location: "us-central1" + vertex_credentials: adroit-crow-413218-a956eef1a2a8.json +``` + +2. Start Proxy + +``` +$ litellm --config /path/to/config.yaml +``` + +3. Make Request use OpenAI Python SDK + +```python +import vertexai + +from vertexai.vision_models import Image, MultiModalEmbeddingModel, Video +from vertexai.vision_models import VideoSegmentConfig +from google.auth.credentials import Credentials + + +LITELLM_PROXY_API_KEY = "sk-1234" +LITELLM_PROXY_BASE = "http://0.0.0.0:4000/vertex-ai" + +import datetime + +class CredentialsWrapper(Credentials): + def __init__(self, token=None): + super().__init__() + self.token = token + self.expiry = None # or set to a future date if needed + + def refresh(self, request): + pass + + def apply(self, headers, token=None): + headers['Authorization'] = f'Bearer {self.token}' + + @property + def expired(self): + return False # Always consider the token as non-expired + + @property + def valid(self): + return True # Always consider the credentials as valid + +credentials = CredentialsWrapper(token=LITELLM_PROXY_API_KEY) + +vertexai.init( + project="adroit-crow-413218", + location="us-central1", + api_endpoint=LITELLM_PROXY_BASE, + credentials = credentials, + api_transport="rest", + +) + +model = MultiModalEmbeddingModel.from_pretrained("multimodalembedding") +image = Image.load_from_file( + "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png" +) + +embeddings = model.get_embeddings( + image=image, + contextual_text="Colosseum", + dimension=1408, +) +print(f"Image Embedding: {embeddings.image_embedding}") +print(f"Text Embedding: {embeddings.text_embedding}") +``` + + + + + +### Text + Image + Video Embeddings + + + + +Text + Image + +```python +response = await litellm.aembedding( + model="vertex_ai/multimodalembedding@001", + input=["hey", "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"] # will be sent as a gcs image +) +``` + +Text + Video + +```python +response = await litellm.aembedding( + model="vertex_ai/multimodalembedding@001", + input=["hey", "gs://my-bucket/embeddings/supermarket-video.mp4"] # will be sent as a gcs image +) +``` + +Image + Video + +```python +response = await litellm.aembedding( + model="vertex_ai/multimodalembedding@001", + input=["gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png", "gs://my-bucket/embeddings/supermarket-video.mp4"] # will be sent as a gcs image +) +``` + + + + + +1. Add model to config.yaml +```yaml +model_list: + - model_name: multimodalembedding@001 + litellm_params: + model: vertex_ai/multimodalembedding@001 + vertex_project: "adroit-crow-413218" + vertex_location: "us-central1" + vertex_credentials: adroit-crow-413218-a956eef1a2a8.json + +litellm_settings: + drop_params: True +``` + +2. Start Proxy + +``` +$ litellm --config /path/to/config.yaml +``` + +3. Make Request use OpenAI Python SDK, Langchain Python SDK + + +Text + Image + +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +# # request sent to model set on litellm proxy, `litellm --model` +response = client.embeddings.create( + model="multimodalembedding@001", + input = ["hey", "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"], +) + +print(response) +``` + +Text + Video +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +# # request sent to model set on litellm proxy, `litellm --model` +response = client.embeddings.create( + model="multimodalembedding@001", + input = ["hey", "gs://my-bucket/embeddings/supermarket-video.mp4"], +) + +print(response) +``` + +Image + Video +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +# # request sent to model set on litellm proxy, `litellm --model` +response = client.embeddings.create( + model="multimodalembedding@001", + input = ["gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png", "gs://my-bucket/embeddings/supermarket-video.mp4"], +) + +print(response) +``` + + + \ No newline at end of file diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index e467711b59d..789cf690285 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -519,6 +519,7 @@ const sidebars = { "providers/vertex_ai/videos", "providers/vertex_partner", "providers/vertex_self_deployed", + "providers/vertex_embedding", "providers/vertex_image", "providers/vertex_batch", "providers/vertex_ocr", From 6341b531a020be1ef6efd5bca4aae4bc11303978 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:46:05 -0700 Subject: [PATCH 11/23] handling BGE URL --- litellm/llms/vertex_ai/common_utils.py | 42 +++++++++++++++++++++++--- 1 file changed, 37 insertions(+), 5 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index aaee922a3f0..430ed909173 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -144,6 +144,36 @@ all_gemini_url_modes = Literal[ ] +def _get_embedding_url( + model: str, + vertex_project: Optional[str], + vertex_location: Optional[str], + vertex_api_version: Literal["v1", "v1beta1"], +) -> Tuple[str, str]: + """ + Get URL for embedding models. + + Handles special patterns: + - bge/endpoint_id -> strips to endpoint_id for endpoints/ routing + - numeric model -> routes to endpoints/ + - regular model -> routes to publishers/google/models/ + """ + from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig + endpoint = "predict" + + # Handle BGE models with pattern bge/endpoint_id (similar to gemma/ pattern) + # After provider split: vertex_ai/bge/123456 -> bge/123456 -> 123456 + if VertexBGEConfig.is_bge_model(model): + model = model.replace("bge/", "", 1) + + 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}" + + return url, endpoint + + def _get_vertex_url( mode: all_gemini_url_modes, model: str, @@ -156,6 +186,7 @@ def _get_vertex_url( endpoint: Optional[str] = None model = litellm.VertexGeminiConfig.get_model_for_vertex_ai_url(model=model) + if mode == "chat": ### SET RUNTIME ENDPOINT ### endpoint = "generateContent" @@ -180,11 +211,12 @@ def _get_vertex_url( if stream is True: url += "?alt=sse" 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}" + return _get_embedding_url( + model=model, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_api_version=vertex_api_version, + ) elif mode == "image_generation": endpoint = "predict" url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}" From 8ea8c674e1fb4a891f92dce5f05d96ae54a9ea3f Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:46:31 -0700 Subject: [PATCH 12/23] fix VertexBGEConfig --- .../llms/vertex_ai/vertex_embeddings/bge.py | 21 +++++++++++++++++-- .../vertex_embeddings/transformation.py | 7 +++++-- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_embeddings/bge.py b/litellm/llms/vertex_ai/vertex_embeddings/bge.py index b8979f55880..2eff0ba96db 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/bge.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/bge.py @@ -4,6 +4,10 @@ Vertex AI BGE (BAAI General Embedding) Configuration BGE models deployed on Vertex AI require different input/output format: - Request: Use "prompt" instead of "content" as the input field - Response: Embeddings are returned directly as arrays, not wrapped in objects + +Model name handling: +- Model names like "bge/endpoint_id" are automatically transformed in common_utils._get_vertex_url() +- This module focuses on request/response transformation only """ from typing import List, Optional, Union @@ -24,6 +28,13 @@ class VertexBGEConfig: BGE (BAAI General Embedding) models use a different request format where the input field is named "prompt" instead of "content". + + Supported model patterns (after provider split in main.py): + - "bge-small-en-v1.5" (model name) + - "bge/204379420394258432" (endpoint ID pattern) + + Note: Model name transformation (bge/ -> numeric ID) is handled automatically + in common_utils._get_vertex_url(). This class focuses on request/response format only. """ @staticmethod @@ -31,13 +42,19 @@ class VertexBGEConfig: """ Check if the model is a BGE (BAAI General Embedding) model. + After provider split in main.py, supports: + - "bge-small-en-v1.5" (model name) + - "bge/204379420394258432" (endpoint ID pattern) + Args: - model: The model name + model: The model name after provider split Returns: bool: True if the model is a BGE model """ - return "bge" in model.lower() + model_lower = model.lower() + # Check for "bge/" prefix (endpoint pattern) or "bge" in model name + return model_lower.startswith("bge/") or "bge" in model_lower @staticmethod def transform_request( diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index 77da3ce7c01..5a3a4a7188a 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -5,7 +5,6 @@ from pydantic import BaseModel from litellm.types.utils import EmbeddingResponse, Usage -from .bge import VertexBGEConfig from .types import * @@ -106,11 +105,12 @@ class VertexAITextEmbeddingConfig(BaseModel): """ Transforms an openai request to a vertex embedding request. """ + # Import here to avoid circular import issues with litellm.__init__ + from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig if model.isdigit(): return self._transform_openai_request_to_fine_tuned_embedding_request( input, optional_params, model ) - if VertexBGEConfig.is_bge_model(model): return VertexBGEConfig.transform_request( input=input, optional_params=optional_params, model=model @@ -216,6 +216,9 @@ class VertexAITextEmbeddingConfig(BaseModel): response, model, model_response ) + # Import here to avoid circular import issues with litellm.__init__ + from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig + if VertexBGEConfig.is_bge_model(model): return VertexBGEConfig.transform_response( response=response, model=model, model_response=model_response From 075a80b7471927541aa082f2c30eb28e816d9d03 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:48:19 -0700 Subject: [PATCH 13/23] test_vertex_ai_bge_with_endpoint_id_pattern --- .../llms/vertex_ai/test_bge_embedding.py | 82 +++++++++++++++++++ 1 file changed, 82 insertions(+) diff --git a/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py b/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py index 636df93b026..75e6f08c822 100644 --- a/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py +++ b/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py @@ -99,3 +99,85 @@ def test_vertex_ai_bge_embedding_with_custom_api_base(): assert len(response.data) == 2 assert "embedding" in response.data[0] + +def test_vertex_ai_bge_with_endpoint_id_pattern(): + """ + Test BGE with vertex_ai/bge/endpoint_id pattern. + + This test verifies that the pattern vertex_ai/bge/204379420394258432 + correctly triggers BGE transformations and routes to the endpoint. + """ + client = HTTPHandler() + + def mock_auth_token(*args, **kwargs): + return "fake-token", "fake-project" + + with patch.object(client, "post") as mock_post, patch( + "litellm.llms.vertex_ai.vertex_embeddings.embedding_handler.VertexEmbedding._ensure_access_token", + side_effect=mock_auth_token + ): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "predictions": [ + [0.1, 0.2, 0.3, 0.4, 0.5], + [0.6, 0.7, 0.8, 0.9, 1.0] + ], + "deployedModelId": "204379420394258432", + "model": "projects/1060139831167/locations/europe-west4/models/baai_bge-base-en", + "modelDisplayName": "baai_bge-base-en", + "modelVersionId": "1" + } + mock_post.return_value = mock_response + + response = litellm.embedding( + model="vertex_ai/bge/204379420394258432", + input=["Hello", "World"], + vertex_project="1060139831167", + vertex_location="europe-west4", + client=client + ) + + mock_post.assert_called_once() + + call_args = mock_post.call_args + kwargs = call_args.kwargs if hasattr(call_args, 'kwargs') else call_args[1] + + if "url" in kwargs: + api_url_called = kwargs["url"] + elif len(call_args[0]) > 0: + api_url_called = call_args[0][0] + else: + api_url_called = "Unknown" + + # Vertex AI may use 'json' or 'data' parameter + if "json" in kwargs: + request_data = kwargs["json"] + elif "data" in kwargs: + request_data = json.loads(kwargs["data"]) + else: + request_data = {} + + print("\n" + "="*50) + print("BGE Endpoint Pattern Test:") + print("="*50) + print(f"Model: vertex_ai/bge/204379420394258432") + print(f"API URL: {api_url_called}") + print("Request Body:") + print(json.dumps(request_data, indent=2)) + print("="*50 + "\n") + + # Verify URL contains the endpoint ID and uses endpoints/ path + assert "204379420394258432" in api_url_called, f"Endpoint ID not in URL: {api_url_called}" + assert "endpoints" in api_url_called, f"Expected 'endpoints' in URL, got: {api_url_called}" + + # Verify BGE-specific request format (uses "prompt" not "content") + assert "instances" in request_data + assert "prompt" in request_data["instances"][0] + assert request_data["instances"][0]["prompt"] == "Hello" + + # Verify response + assert isinstance(response.data, list) + assert len(response.data) == 2 + + From b7fe25c97db1a3a77b1bc23b65c43e964cbda42a Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:53:32 -0700 Subject: [PATCH 14/23] docs vertex BGE --- .../docs/providers/vertex_embedding.md | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/docs/my-website/docs/providers/vertex_embedding.md b/docs/my-website/docs/providers/vertex_embedding.md index 25580935387..023db6130f7 100644 --- a/docs/my-website/docs/providers/vertex_embedding.md +++ b/docs/my-website/docs/providers/vertex_embedding.md @@ -179,6 +179,70 @@ print(response) +## **BGE Embeddings** + +Use BGE (Baidu General Embedding) models deployed on Vertex AI. + +### Usage + + + + +```python showLineNumbers title="Using BGE on Vertex AI" +import litellm + +response = litellm.embedding( + model="vertex_ai/bge/", + input=["Hello", "World"], + vertex_project="your-project-id", + vertex_location="your-location" +) + +print(response) +``` + + + + + +1. Add model to config.yaml +```yaml showLineNumbers title="config.yaml" +model_list: + - model_name: bge-embedding + litellm_params: + model: vertex_ai/bge/ + vertex_project: "your-project-id" + vertex_location: "us-central1" + vertex_credentials: your-credentials.json + +litellm_settings: + drop_params: True +``` + +2. Start Proxy + +```bash +$ litellm --config /path/to/config.yaml +``` + +3. Make Request using OpenAI Python SDK + +```python showLineNumbers title="Making requests to BGE" +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +response = client.embeddings.create( + model="bge-embedding", + input=["good morning from litellm", "this is another item"] +) + +print(response) +``` + + + + ## **Multi-Modal Embeddings** From a79002c1fe42c204c1f9c1e3f15b009e80744d3f Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:57:20 -0700 Subject: [PATCH 15/23] docs --- docs/my-website/docs/providers/vertex_embedding.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/docs/my-website/docs/providers/vertex_embedding.md b/docs/my-website/docs/providers/vertex_embedding.md index 023db6130f7..ad2c03debb1 100644 --- a/docs/my-website/docs/providers/vertex_embedding.md +++ b/docs/my-website/docs/providers/vertex_embedding.md @@ -240,6 +240,18 @@ response = client.embeddings.create( print(response) ``` +Using a Private Service Connect (PSC) endpoint + +```yaml showLineNumbers title="config.yaml (PSC)" +model_list: + - model_name: bge-small-en-v1.5 + litellm_params: + model: vertex_ai/1234567890 + api_base: http://10.96.32.8 # Your PSC IP + vertex_project: my-project-id #optional + vertex_location: us-central1 #optional +``` + From fcc108b554867de286ac56fe4d66f25c96fbdb1d Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:57:35 -0700 Subject: [PATCH 16/23] docs fix --- docs/my-website/docs/providers/vertex_embedding.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/my-website/docs/providers/vertex_embedding.md b/docs/my-website/docs/providers/vertex_embedding.md index ad2c03debb1..5656ade337b 100644 --- a/docs/my-website/docs/providers/vertex_embedding.md +++ b/docs/my-website/docs/providers/vertex_embedding.md @@ -246,7 +246,7 @@ Using a Private Service Connect (PSC) endpoint model_list: - model_name: bge-small-en-v1.5 litellm_params: - model: vertex_ai/1234567890 + model: vertex_ai/bge/1234567890 api_base: http://10.96.32.8 # Your PSC IP vertex_project: my-project-id #optional vertex_location: us-central1 #optional From c0a083ff61546bf0aeca3ea4e952d0281508382e Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 29 Oct 2025 09:53:23 -0700 Subject: [PATCH 17/23] fix VertexAIModelRoute --- litellm/llms/vertex_ai/common_utils.py | 42 +++----------------------- 1 file changed, 5 insertions(+), 37 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 430ed909173..aaee922a3f0 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -144,36 +144,6 @@ all_gemini_url_modes = Literal[ ] -def _get_embedding_url( - model: str, - vertex_project: Optional[str], - vertex_location: Optional[str], - vertex_api_version: Literal["v1", "v1beta1"], -) -> Tuple[str, str]: - """ - Get URL for embedding models. - - Handles special patterns: - - bge/endpoint_id -> strips to endpoint_id for endpoints/ routing - - numeric model -> routes to endpoints/ - - regular model -> routes to publishers/google/models/ - """ - from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig - endpoint = "predict" - - # Handle BGE models with pattern bge/endpoint_id (similar to gemma/ pattern) - # After provider split: vertex_ai/bge/123456 -> bge/123456 -> 123456 - if VertexBGEConfig.is_bge_model(model): - model = model.replace("bge/", "", 1) - - 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}" - - return url, endpoint - - def _get_vertex_url( mode: all_gemini_url_modes, model: str, @@ -186,7 +156,6 @@ def _get_vertex_url( endpoint: Optional[str] = None model = litellm.VertexGeminiConfig.get_model_for_vertex_ai_url(model=model) - if mode == "chat": ### SET RUNTIME ENDPOINT ### endpoint = "generateContent" @@ -211,12 +180,11 @@ def _get_vertex_url( if stream is True: url += "?alt=sse" elif mode == "embedding": - return _get_embedding_url( - model=model, - vertex_project=vertex_project, - vertex_location=vertex_location, - vertex_api_version=vertex_api_version, - ) + 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}" elif mode == "image_generation": endpoint = "predict" url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}" From 8a9c9af55f23d67b80450aeaaaf36c8a0d80a097 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 29 Oct 2025 09:53:52 -0700 Subject: [PATCH 18/23] from ..common_utils import VertexAIError, get_vertex_base_model_name add --- litellm/llms/vertex_ai/vertex_model_garden/main.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_model_garden/main.py b/litellm/llms/vertex_ai/vertex_model_garden/main.py index 225e75a5add..fe7d0862e02 100644 --- a/litellm/llms/vertex_ai/vertex_model_garden/main.py +++ b/litellm/llms/vertex_ai/vertex_model_garden/main.py @@ -22,7 +22,7 @@ import httpx # type: ignore from litellm.utils import ModelResponse -from ..common_utils import VertexAIError +from ..common_utils import VertexAIError, get_vertex_base_model_name from ..vertex_llm_base import VertexBase @@ -89,7 +89,7 @@ class VertexAIModelGardenModels(VertexBase): message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""", ) try: - model = model.replace("openai/", "") + model = get_vertex_base_model_name(model=model) vertex_httpx_logic = VertexLLM() access_token, project_id = vertex_httpx_logic._ensure_access_token( From 87b75afe12d6b2bc4644979e596bff5b750ea2dd Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 29 Oct 2025 09:54:20 -0700 Subject: [PATCH 19/23] fix VertexAIGemmaModels --- litellm/llms/vertex_ai/vertex_gemma_models/main.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/main.py b/litellm/llms/vertex_ai/vertex_gemma_models/main.py index 8203b285ebd..41bd6b5431e 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/main.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/main.py @@ -25,7 +25,7 @@ import httpx # type: ignore from litellm.utils import ModelResponse -from ..common_utils import VertexAIError +from ..common_utils import VertexAIError, get_vertex_base_model_name from ..vertex_llm_base import VertexBase @@ -82,7 +82,8 @@ class VertexAIGemmaModels(VertexBase): message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""", ) try: - model = model.replace("gemma/", "") + + model = get_vertex_base_model_name(model=model) vertex_httpx_logic = VertexLLM() access_token, project_id = vertex_httpx_logic._ensure_access_token( From bfa7f12d4c78a7e89222815c51d3bc517a7f9c3d Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 29 Oct 2025 09:54:59 -0700 Subject: [PATCH 20/23] fix get_vertex_base_model_name --- litellm/llms/vertex_ai/vertex_llm_base.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index a5c44617fab..ce50bf311e1 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -19,6 +19,7 @@ from .common_utils import ( _get_gemini_url, _get_vertex_url, all_gemini_url_modes, + get_vertex_base_model_name, is_global_only_vertex_model, ) @@ -327,10 +328,13 @@ class VertexBase: # Check if this is a PSC endpoint or custom deployment # PSC/custom endpoints need the full path structure if vertex_project and vertex_location and model: + # Strip routing prefixes (bge/, gemma/, etc.) for endpoint URL construction + model_for_url = get_vertex_base_model_name(model=model) + # Check if model is numeric (endpoint ID) or if api_base doesn't contain googleapis.com # These are indicators of PSC/custom endpoints is_psc_or_custom = ( - "googleapis.com" not in api_base.lower() or model.isdigit() + "googleapis.com" not in api_base.lower() or model_for_url.isdigit() ) if is_psc_or_custom: @@ -342,7 +346,7 @@ class VertexBase: version, vertex_project, vertex_location, - model, + model_for_url, endpoint, ) else: From fe03833d3bae772ffd2d05603d1348186a9a874a Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 29 Oct 2025 10:01:06 -0700 Subject: [PATCH 21/23] test_vertex_ai_bge_psc_endpoint_url_construction --- .../llms/vertex_ai/test_bge_embedding.py | 68 +++++++++++++++++++ 1 file changed, 68 insertions(+) diff --git a/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py b/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py index 75e6f08c822..156ab95184a 100644 --- a/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py +++ b/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py @@ -181,3 +181,71 @@ def test_vertex_ai_bge_with_endpoint_id_pattern(): assert len(response.data) == 2 +def test_vertex_ai_bge_psc_endpoint_url_construction(): + """ + Test that BGE models with PSC endpoints construct correct URL without bge/ prefix. + + Verifies that vertex_ai/bge/378943383978115072 with api_base http://10.128.16.2 + constructs URL: http://10.128.16.2/v1/projects/{project}/locations/{location}/endpoints/378943383978115072:predict + + The bge/ prefix should be stripped from the endpoint URL. + """ + client = HTTPHandler() + + def mock_auth_token(*args, **kwargs): + return "fake-token", "gen-lang-client-0682925754" + + with patch.object(client, "post") as mock_post, patch( + "litellm.llms.vertex_ai.vertex_embeddings.embedding_handler.VertexEmbedding._ensure_access_token", + side_effect=mock_auth_token + ): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "predictions": [ + [0.1, 0.2, 0.3, 0.4, 0.5] + ] + } + mock_post.return_value = mock_response + + response = litellm.embedding( + model="vertex_ai/bge/378943383978115072", + input=["The food was delicious and the waiter.."], + api_base="http://10.128.16.2", + vertex_project="gen-lang-client-0682925754", + vertex_location="us-central1", + client=client + ) + + mock_post.assert_called_once() + + call_args = mock_post.call_args + kwargs = call_args.kwargs if hasattr(call_args, 'kwargs') else call_args[1] + + if "url" in kwargs: + api_url_called = kwargs["url"] + elif len(call_args[0]) > 0: + api_url_called = call_args[0][0] + else: + api_url_called = "Unknown" + + print("\n" + "="*50) + print("PSC Endpoint URL Construction Test:") + print("="*50) + print(f"Model: vertex_ai/bge/378943383978115072") + print(f"API Base: http://10.128.16.2") + print(f"Constructed URL: {api_url_called}") + print("="*50 + "\n") + + # Verify the URL is constructed correctly + expected_url = "http://10.128.16.2/v1/projects/gen-lang-client-0682925754/locations/us-central1/endpoints/378943383978115072:predict" + assert api_url_called == expected_url, f"Expected URL: {expected_url}, Got: {api_url_called}" + + # Verify bge/ prefix is NOT in the URL + assert "bge/" not in api_url_called, f"URL should not contain 'bge/' prefix: {api_url_called}" + + # Verify response works + assert isinstance(response.data, list) + assert len(response.data) == 1 + + From 2201e12accfd27e531937d2f3a1b00dd400a9fe1 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 2 Dec 2025 22:08:23 +0530 Subject: [PATCH 22/23] Fix import error --- litellm/llms/vertex_ai/common_utils.py | 87 +++++++++++++++++-- .../test_vertex_ai_psc_endpoint_support.py | 5 +- 2 files changed, 83 insertions(+), 9 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index aaee922a3f0..836234f6f13 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -31,9 +31,11 @@ class VertexAIModelRoute(str, Enum): PARTNER_MODELS = "partner_models" GEMINI = "gemini" GEMMA = "gemma" + BGE = "bge" MODEL_GARDEN = "model_garden" NON_GEMINI = "non_gemini" +VERTEX_AI_MODEL_ROUTES = [f"{route.value}/" for route in VertexAIModelRoute] def get_vertex_ai_model_route( model: str, litellm_params: Optional[dict] = None @@ -81,7 +83,11 @@ def get_vertex_ai_model_route( # Check for partner models (llama, mistral, claude, etc.) if VertexAIPartnerModels.is_vertex_partner_model(model=model): return VertexAIModelRoute.PARTNER_MODELS - + + # Check for BGE models + if "bge/" in model or "bge" in model.lower(): + return VertexAIModelRoute.BGE + # Check for gemma models if "gemma/" in model: return VertexAIModelRoute.GEMMA @@ -144,6 +150,71 @@ all_gemini_url_modes = Literal[ ] +def get_vertex_base_model_name(model: str) -> str: + """ + Strip routing prefixes from model name for PSC/endpoint URL construction. + + Patterns like "bge/", "gemma/", "openai/" are used for internal routing but + should not appear in the actual endpoint URL. Routing prefixes are derived + from VertexAIModelRoute enum values. + + Args: + model: The model name with potential prefix (e.g., "bge/123456", "gemma/gemma-3-12b-it") + + Returns: + str: The model name without routing prefix (e.g., "123456", "gemma-3-12b-it") + + Examples: + >>> get_vertex_base_model_name("bge/378943383978115072") + "378943383978115072" + + >>> get_vertex_base_model_name("gemma/gemma-3-12b-it") + "gemma-3-12b-it" + + >>> get_vertex_base_model_name("openai/gpt-oss-120b") + "gpt-oss-120b" + + >>> get_vertex_base_model_name("1234567890") + "1234567890" + """ + # Derive routing prefixes from VertexAIModelRoute enum + # Map specific routes to their prefixes (some routes like PARTNER_MODELS, GEMINI don't have prefixes) + + + for route in VERTEX_AI_MODEL_ROUTES: + if model.startswith(route): + return model.replace(route, "", 1) + + return model + + +def _get_embedding_url( + model: str, + vertex_project: Optional[str], + vertex_location: Optional[str], + vertex_api_version: Literal["v1", "v1beta1"], +) -> Tuple[str, str]: + """ + Get URL for embedding models. + + Handles special patterns: + - bge/endpoint_id -> strips to endpoint_id for endpoints/ routing + - numeric model -> routes to endpoints/ + - regular model -> routes to publishers/google/models/ + """ + endpoint = "predict" + + # Strip routing prefixes (bge/, gemma/, etc.) for endpoint URL construction + model = get_vertex_base_model_name(model=model) + + 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}" + + return url, endpoint + + def _get_vertex_url( mode: all_gemini_url_modes, model: str, @@ -156,6 +227,7 @@ def _get_vertex_url( endpoint: Optional[str] = None model = litellm.VertexGeminiConfig.get_model_for_vertex_ai_url(model=model) + if mode == "chat": ### SET RUNTIME ENDPOINT ### endpoint = "generateContent" @@ -180,11 +252,12 @@ def _get_vertex_url( if stream is True: url += "?alt=sse" 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}" + return _get_embedding_url( + model=model, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_api_version=vertex_api_version, + ) elif mode == "image_generation": endpoint = "predict" url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}" @@ -870,4 +943,4 @@ class VertexAITokenCounter(BaseTokenCounter): original_response=result, ) - return None + return None \ No newline at end of file diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py index 46f365094c0..c158c93be9d 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py @@ -5,9 +5,10 @@ Tests that LiteLLM properly constructs URLs when using custom api_base for PSC endpoints. """ -import pytest -import sys import os +import sys + +import pytest # Add the litellm package to the path sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../../..")) From 46ebf425d56b6369f61188b64e26c8daad87a373 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 4 Dec 2025 21:39:42 +0530 Subject: [PATCH 23/23] Fix : test_vertexai_model_garden_model_completion --- litellm/llms/vertex_ai/common_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 836234f6f13..c0dfda00abe 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -34,6 +34,7 @@ class VertexAIModelRoute(str, Enum): BGE = "bge" MODEL_GARDEN = "model_garden" NON_GEMINI = "non_gemini" + OPENAI_COMPATIBLE = "openai" VERTEX_AI_MODEL_ROUTES = [f"{route.value}/" for route in VertexAIModelRoute] @@ -179,8 +180,6 @@ def get_vertex_base_model_name(model: str) -> str: """ # Derive routing prefixes from VertexAIModelRoute enum # Map specific routes to their prefixes (some routes like PARTNER_MODELS, GEMINI don't have prefixes) - - for route in VERTEX_AI_MODEL_ROUTES: if model.startswith(route): return model.replace(route, "", 1)