tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-23 20:26:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

Revert "[Feat] VertexAI - Add BGE Embeddings support (#16033)" (#16677)

Browse Source 
This reverts commit 7133488282.
This commit is contained in:
Ishaan Jaff
2025-11-14 17:41:06 -08:00
committed by GitHub
parent f8c022c8f5
commit 8a43fbe8f7
15 changed files with 526 additions and 1621 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/providers/vertex.md
+509 -47
View File
@@ -1604,53 +1604,6 @@ 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
@@ -2089,6 +2042,515 @@ curl http://0.0.0.0:4000/v1/chat/completions \
| code-gecko@latest| `completion('code-gecko@latest', messages)` |
## **Embedding Models**
#### Usage - Embedding
<Tabs>
<TabItem value="sdk" label="SDK">
```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)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY">
1. Add model to config.yaml
```yaml
model_list:
- model_name: snowflake-arctic-embed-m-long-1731622468876
litellm_params:
model: vertex_ai/<your-model-id>
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)
```
</TabItem>
</Tabs>
#### 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/<your-model-id>", 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
<Tabs>
<TabItem value="sdk" label="SDK">
```python
response = litellm.embedding(
model="vertex_ai/text-embedding-004",
input=["good morning from litellm", "gm"]
input_type = "RETRIEVAL_DOCUMENT",
dimensions=1,
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY">
```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)
```
</TabItem>
</Tabs>
### 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)
<Tabs>
<TabItem value="sdk" label="SDK">
```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,
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY">
```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)
```
</TabItem>
</Tabs>
## **Multi-Modal Embeddings**
Known Limitations:
- Only supports 1 image / video / image per request
- Only supports GCS or base64 encoded images / videos
### Usage
<Tabs>
<TabItem value="sdk" label="SDK">
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
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY (Unified Endpoint)">
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
<Tabs>
<TabItem value="OpenAI SDK" label="OpenAI 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)
```
</TabItem>
<TabItem value="langchain" label="Langchain">
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)
```
</TabItem>
</Tabs>
</TabItem>
<TabItem value="proxy-vtx" label="LiteLLM PROXY (Vertex SDK)">
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}")
```
</TabItem>
</Tabs>
### Text + Image + Video Embeddings
<Tabs>
<TabItem value="sdk" label="SDK">
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
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY (Unified Endpoint)">
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)
```
</TabItem>
</Tabs>
## **Gemini TTS (Text-to-Speech) Audio Output**
:::info
docs/my-website/docs/providers/vertex_embedding.md
-587
View File
@@ -1,587 +0,0 @@
import Image from '@theme/IdealImage';
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Vertex AI Embedding
## Usage - Embedding
<Tabs>
<TabItem value="sdk" label="SDK">
```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)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY">
1. Add model to config.yaml
```yaml
model_list:
- model_name: snowflake-arctic-embed-m-long-1731622468876
litellm_params:
model: vertex_ai/<your-model-id>
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)
```
</TabItem>
</Tabs>
#### 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/<your-model-id>", 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
<Tabs>
<TabItem value="sdk" label="SDK">
```python
response = litellm.embedding(
model="vertex_ai/text-embedding-004",
input=["good morning from litellm", "gm"]
input_type = "RETRIEVAL_DOCUMENT",
dimensions=1,
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY">
```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)
```
</TabItem>
</Tabs>
### 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)
<Tabs>
<TabItem value="sdk" label="SDK">
```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,
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY">
```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)
```
</TabItem>
</Tabs>
## **BGE Embeddings**
Use BGE (Baidu General Embedding) models deployed on Vertex AI.
### Usage
<Tabs>
<TabItem value="sdk" label="SDK">
```python showLineNumbers title="Using BGE on Vertex AI"
import litellm
response = litellm.embedding(
model="vertex_ai/bge/<your-endpoint-id>",
input=["Hello", "World"],
vertex_project="your-project-id",
vertex_location="your-location"
)
print(response)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY">
1. Add model to config.yaml
```yaml showLineNumbers title="config.yaml"
model_list:
- model_name: bge-embedding
litellm_params:
model: vertex_ai/bge/<your-endpoint-id>
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
```
</TabItem>
</Tabs>
## **Multi-Modal Embeddings**
Known Limitations:
- Only supports 1 image / video / image per request
- Only supports GCS or base64 encoded images / videos
### Usage
<Tabs>
<TabItem value="sdk" label="SDK">
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
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY (Unified Endpoint)">
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
<Tabs>
<TabItem value="OpenAI SDK" label="OpenAI 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)
```
</TabItem>
<TabItem value="langchain" label="Langchain">
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)
```
</TabItem>
</Tabs>
</TabItem>
<TabItem value="proxy-vtx" label="LiteLLM PROXY (Vertex SDK)">
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}")
```
</TabItem>
</Tabs>
### Text + Image + Video Embeddings
<Tabs>
<TabItem value="sdk" label="SDK">
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
)
```
</TabItem>
<TabItem value="proxy" label="LiteLLM PROXY (Unified Endpoint)">
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)
```
</TabItem>
</Tabs>
docs/my-website/sidebars.js
-1
View File
@@ -497,7 +497,6 @@ 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",
litellm/llms/vertex_ai/batches/handler.py
-8
View File
@@ -61,10 +61,6 @@ 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 = {
@@ -170,10 +166,6 @@ 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 = {
litellm/llms/vertex_ai/common_utils.py
+7 -88
View File
@@ -31,11 +31,9 @@ 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
@@ -62,9 +60,6 @@ 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,
@@ -74,20 +69,11 @@ 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
@@ -150,71 +136,6 @@ 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,
@@ -227,7 +148,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"
@@ -252,12 +172,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}"
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py
-4
View File
@@ -79,10 +79,6 @@ 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(
litellm/llms/vertex_ai/vertex_embeddings/bge.py
-182
View File
@@ -1,182 +0,0 @@
"""
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
litellm/llms/vertex_ai/vertex_embeddings/transformation.py
+2 -19
View File
@@ -105,16 +105,10 @@ 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] = []
@@ -173,9 +167,6 @@ 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
@@ -192,8 +183,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.
@@ -215,14 +206,6 @@ 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"]
litellm/llms/vertex_ai/vertex_embeddings/types.py
+1 -7
View File
@@ -25,12 +25,6 @@ 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):
@@ -50,7 +44,7 @@ class EmbeddingParameters(TypedDict, total=False):
class VertexEmbeddingRequest(TypedDict, total=False):
instances: Union[List[TextEmbeddingInput], List[TextEmbeddingBGEInput], List[TextEmbeddingFineTunedInput]]
instances: Union[List[TextEmbeddingInput], List[TextEmbeddingFineTunedInput]]
parameters: Optional[Union[EmbeddingParameters, TextEmbeddingFineTunedParameters]]
litellm/llms/vertex_ai/vertex_gemma_models/main.py
+2 -3
View File
@@ -25,7 +25,7 @@ import httpx # type: ignore
from litellm.utils import ModelResponse
from ..common_utils import VertexAIError, get_vertex_base_model_name
from ..common_utils import VertexAIError
from ..vertex_llm_base import VertexBase
@@ -82,8 +82,7 @@ class VertexAIGemmaModels(VertexBase):
message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""",
)
try:
model = get_vertex_base_model_name(model=model)
model = model.replace("gemma/", "")
vertex_httpx_logic = VertexLLM()
access_token, project_id = vertex_httpx_logic._ensure_access_token(
litellm/llms/vertex_ai/vertex_llm_base.py
+3 -49
View File
@@ -19,7 +19,6 @@ from .common_utils import (
_get_gemini_url,
_get_vertex_url,
all_gemini_url_modes,
get_vertex_base_model_name,
is_global_only_vertex_model,
)
@@ -242,9 +241,6 @@ 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
@@ -293,18 +289,9 @@ 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]
@@ -324,37 +311,8 @@ class VertexBase:
if gemini_api_key is not None:
auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment]
else:
# 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)
url = "{}:{}".format(api_base, endpoint)
if stream is True:
url = url + "?alt=sse"
return auth_header, url
@@ -381,7 +339,6 @@ class VertexBase:
Returns
token, url
"""
version: Optional[Literal["v1beta1", "v1"]] = None
if custom_llm_provider == "gemini":
url, endpoint = _get_gemini_url(
mode=mode,
@@ -397,7 +354,7 @@ class VertexBase:
)
### SET RUNTIME ENDPOINT ###
version = (
version: Literal["v1beta1", "v1"] = (
"v1beta1" if should_use_v1beta1_features is True else "v1"
)
url, endpoint = _get_vertex_url(
@@ -418,9 +375,6 @@ class VertexBase:
stream=stream,
url=url,
model=model,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_api_version=version,
)
def _handle_reauthentication(
litellm/llms/vertex_ai/vertex_model_garden/main.py
+2 -6
View File
@@ -22,7 +22,7 @@ import httpx # type: ignore
from litellm.utils import ModelResponse
from ..common_utils import VertexAIError, get_vertex_base_model_name
from ..common_utils import VertexAIError
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 = get_vertex_base_model_name(model=model)
model = model.replace("openai/", "")
vertex_httpx_logic = VertexLLM()
access_token, project_id = vertex_httpx_logic._ensure_access_token(
@@ -123,10 +123,6 @@ 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(
tests/test_litellm/llms/vertex_ai/test_bge_embedding.py
-251
View File
@@ -1,251 +0,0 @@
"""
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
tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py
-111
View File
@@ -1,111 +0,0 @@
"""
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
)
tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py
-258
View File
@@ -1,258 +0,0 @@
"""
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}"