diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md
index da2997f6202..7b762b59560 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/docs/my-website/docs/providers/vertex_embedding.md b/docs/my-website/docs/providers/vertex_embedding.md
new file mode 100644
index 00000000000..5656ade337b
--- /dev/null
+++ b/docs/my-website/docs/providers/vertex_embedding.md
@@ -0,0 +1,587 @@
+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)
+```
+
+
+
+## **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)
+```
+
+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/bge/1234567890
+ api_base: http://10.96.32.8 # Your PSC IP
+ vertex_project: my-project-id #optional
+ vertex_location: us-central1 #optional
+```
+
+
+
+
+## **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 cebe31a8e11..b7143490dea 100644
--- a/docs/my-website/sidebars.js
+++ b/docs/my-website/sidebars.js
@@ -544,6 +544,7 @@ const sidebars = {
"providers/vertex_ai/videos",
"providers/vertex_partner",
"providers/vertex_self_deployed",
+ "providers/vertex_embedding",
"providers/vertex_image",
"providers/vertex_speech",
"providers/vertex_batch",
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 a3ea8afe40d..789c880122f 100644
--- a/litellm/llms/vertex_ai/common_utils.py
+++ b/litellm/llms/vertex_ai/common_utils.py
@@ -32,9 +32,12 @@ class VertexAIModelRoute(str, Enum):
PARTNER_MODELS = "partner_models"
GEMINI = "gemini"
GEMMA = "gemma"
+ BGE = "bge"
MODEL_GARDEN = "model_garden"
NON_GEMINI = "non_gemini"
+ OPENAI_COMPATIBLE = "openai"
+VERTEX_AI_MODEL_ROUTES = [f"{route.value}/" for route in VertexAIModelRoute]
def get_vertex_ai_model_route(
model: str, litellm_params: Optional[dict] = None
@@ -61,6 +64,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,
@@ -70,11 +76,20 @@ 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
-
+
+ # 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
@@ -137,6 +152,69 @@ 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,
@@ -149,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"
@@ -173,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}"
@@ -863,4 +943,4 @@ class VertexAITokenCounter(BaseTokenCounter):
original_response=result,
)
- return None
+ return None
\ No newline at end of file
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/bge.py b/litellm/llms/vertex_ai/vertex_embeddings/bge.py
new file mode 100644
index 00000000000..2eff0ba96db
--- /dev/null
+++ b/litellm/llms/vertex_ai/vertex_embeddings/bge.py
@@ -0,0 +1,182 @@
+"""
+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
+
+from litellm.types.utils import EmbeddingResponse, Usage
+
+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".
+
+ 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
+ def is_bge_model(model: str) -> bool:
+ """
+ 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 after provider split
+
+ Returns:
+ bool: True if the model is a BGE model
+ """
+ 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(
+ 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
+
+ @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
+
+ 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",
+ "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 97af558041d..5a3a4a7188a 100644
--- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py
+++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py
@@ -105,10 +105,16 @@ 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
+ )
vertex_request: VertexEmbeddingRequest = VertexEmbeddingRequest()
vertex_text_embedding_input_list: List[TextEmbeddingInput] = []
@@ -167,6 +173,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
@@ -183,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.
@@ -206,6 +215,14 @@ class VertexAITextEmbeddingConfig(BaseModel):
return self._transform_vertex_response_to_openai_for_fine_tuned_models(
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
+ )
_predictions = response["predictions"]
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]]
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(
diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py
index 9ddbc461a70..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,
)
@@ -241,6 +242,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 +293,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 +324,37 @@ 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:
+ # 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_for_url.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_for_url,
+ 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 +381,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 +397,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 +418,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..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(
@@ -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_bge_embedding.py b/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py
new file mode 100644
index 00000000000..156ab95184a
--- /dev/null
+++ b/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py
@@ -0,0 +1,251 @@
+"""
+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
+ # BGE models return embeddings directly as arrays, not wrapped in objects
+ 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": "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
+
+ 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]
+
+
+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
+
+
+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
+
+
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..20150501adf
--- /dev/null
+++ b/tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py
@@ -0,0 +1,111 @@
+"""
+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_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.
+
+ 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
+ )
+
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..c158c93be9d
--- /dev/null
+++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py
@@ -0,0 +1,259 @@
+"""
+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 os
+import sys
+
+import pytest
+
+# 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}"
+