From 8a43fbe8f7bdba1ff2151a9e490c32c75cab6267 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 14 Nov 2025 17:41:06 -0800 Subject: [PATCH] Revert "[Feat] VertexAI - Add BGE Embeddings support (#16033)" (#16677) This reverts commit 713348828213493407dde134c0f93bfc4fecea67. --- docs/my-website/docs/providers/vertex.md | 556 +++++++++++++++-- .../docs/providers/vertex_embedding.md | 587 ------------------ docs/my-website/sidebars.js | 1 - litellm/llms/vertex_ai/batches/handler.py | 8 - litellm/llms/vertex_ai/common_utils.py | 95 +-- .../vertex_ai_context_caching.py | 4 - .../llms/vertex_ai/vertex_embeddings/bge.py | 182 ------ .../vertex_embeddings/transformation.py | 21 +- .../llms/vertex_ai/vertex_embeddings/types.py | 8 +- .../vertex_ai/vertex_gemma_models/main.py | 5 +- litellm/llms/vertex_ai/vertex_llm_base.py | 52 +- .../vertex_ai/vertex_model_garden/main.py | 8 +- .../llms/vertex_ai/test_bge_embedding.py | 251 -------- .../test_bge_response_transformation.py | 111 ---- .../test_vertex_ai_psc_endpoint_support.py | 258 -------- 15 files changed, 526 insertions(+), 1621 deletions(-) delete mode 100644 docs/my-website/docs/providers/vertex_embedding.md delete mode 100644 litellm/llms/vertex_ai/vertex_embeddings/bge.py delete mode 100644 tests/test_litellm/llms/vertex_ai/test_bge_embedding.py delete mode 100644 tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py delete mode 100644 tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 4d7e85f388..874b637e4d 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -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 + + + + +```python +import litellm +from litellm import embedding +litellm.vertex_project = "hardy-device-38811" # Your Project ID +litellm.vertex_location = "us-central1" # proj location + +response = embedding( + model="vertex_ai/textembedding-gecko", + input=["good morning from litellm"], +) +print(response) +``` + + + + + +1. Add model to config.yaml +```yaml +model_list: + - model_name: snowflake-arctic-embed-m-long-1731622468876 + litellm_params: + model: vertex_ai/ + vertex_project: "adroit-crow-413218" + vertex_location: "us-central1" + vertex_credentials: adroit-crow-413218-a956eef1a2a8.json + +litellm_settings: + drop_params: True +``` + +2. Start Proxy + +``` +$ litellm --config /path/to/config.yaml +``` + +3. Make Request using OpenAI Python SDK, Langchain Python SDK + +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +response = client.embeddings.create( + model="snowflake-arctic-embed-m-long-1731622468876", + input = ["good morning from litellm", "this is another item"], +) + +print(response) +``` + + + + + +#### Supported Embedding Models +All models listed [here](https://github.com/BerriAI/litellm/blob/57f37f743886a0249f630a6792d49dffc2c5d9b7/model_prices_and_context_window.json#L835) are supported + +| Model Name | Function Call | +|--------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| text-embedding-004 | `embedding(model="vertex_ai/text-embedding-004", input)` | +| text-multilingual-embedding-002 | `embedding(model="vertex_ai/text-multilingual-embedding-002", input)` | +| textembedding-gecko | `embedding(model="vertex_ai/textembedding-gecko", input)` | +| textembedding-gecko-multilingual | `embedding(model="vertex_ai/textembedding-gecko-multilingual", input)` | +| textembedding-gecko-multilingual@001 | `embedding(model="vertex_ai/textembedding-gecko-multilingual@001", input)` | +| textembedding-gecko@001 | `embedding(model="vertex_ai/textembedding-gecko@001", input)` | +| textembedding-gecko@003 | `embedding(model="vertex_ai/textembedding-gecko@003", input)` | +| text-embedding-preview-0409 | `embedding(model="vertex_ai/text-embedding-preview-0409", input)` | +| text-multilingual-embedding-preview-0409 | `embedding(model="vertex_ai/text-multilingual-embedding-preview-0409", input)` | +| Fine-tuned OR Custom Embedding models | `embedding(model="vertex_ai/", input)` | + +### Supported OpenAI (Unified) Params + +| [param](../embedding/supported_embedding.md#input-params-for-litellmembedding) | type | [vertex equivalent](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api) | +|-------|-------------|--------------------| +| `input` | **string or List[string]** | `instances` | +| `dimensions` | **int** | `output_dimensionality` | +| `input_type` | **Literal["RETRIEVAL_QUERY","RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING", "FACT_VERIFICATION"]** | `task_type` | + +#### Usage with OpenAI (Unified) Params + + + + + +```python +response = litellm.embedding( + model="vertex_ai/text-embedding-004", + input=["good morning from litellm", "gm"] + input_type = "RETRIEVAL_DOCUMENT", + dimensions=1, +) +``` + + + + +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +response = client.embeddings.create( + model="text-embedding-004", + input = ["good morning from litellm", "gm"], + dimensions=1, + extra_body = { + "input_type": "RETRIEVAL_QUERY", + } +) + +print(response) +``` + + + + +### Supported Vertex Specific Params + +| param | type | +|-------|-------------| +| `auto_truncate` | **bool** | +| `task_type` | **Literal["RETRIEVAL_QUERY","RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY", "CLASSIFICATION", "CLUSTERING", "QUESTION_ANSWERING", "FACT_VERIFICATION"]** | +| `title` | **str** | + +#### Usage with Vertex Specific Params (Use `task_type` and `title`) + +You can pass any vertex specific params to the embedding model. Just pass them to the embedding function like this: + +[Relevant Vertex AI doc with all embedding params](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api#request_body) + + + + +```python +response = litellm.embedding( + model="vertex_ai/text-embedding-004", + input=["good morning from litellm", "gm"] + task_type = "RETRIEVAL_DOCUMENT", + title = "test", + dimensions=1, + auto_truncate=True, +) +``` + + + + +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +response = client.embeddings.create( + model="text-embedding-004", + input = ["good morning from litellm", "gm"], + dimensions=1, + extra_body = { + "task_type": "RETRIEVAL_QUERY", + "auto_truncate": True, + "title": "test", + } +) + +print(response) +``` + + + +## **Multi-Modal Embeddings** + + +Known Limitations: +- Only supports 1 image / video / image per request +- Only supports GCS or base64 encoded images / videos + +### Usage + + + + +Using GCS Images + +```python +response = await litellm.aembedding( + model="vertex_ai/multimodalembedding@001", + input="gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png" # will be sent as a gcs image +) +``` + +Using base 64 encoded images + +```python +response = await litellm.aembedding( + model="vertex_ai/multimodalembedding@001", + input="data:image/jpeg;base64,..." # will be sent as a base64 encoded image +) +``` + + + + +1. Add model to config.yaml +```yaml +model_list: + - model_name: multimodalembedding@001 + litellm_params: + model: vertex_ai/multimodalembedding@001 + vertex_project: "adroit-crow-413218" + vertex_location: "us-central1" + vertex_credentials: adroit-crow-413218-a956eef1a2a8.json + +litellm_settings: + drop_params: True +``` + +2. Start Proxy + +``` +$ litellm --config /path/to/config.yaml +``` + +3. Make Request use OpenAI Python SDK, Langchain Python SDK + + + + + + +Requests with GCS Image / Video URI + +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +# # request sent to model set on litellm proxy, `litellm --model` +response = client.embeddings.create( + model="multimodalembedding@001", + input = "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png", +) + +print(response) +``` + +Requests with base64 encoded images + +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +# # request sent to model set on litellm proxy, `litellm --model` +response = client.embeddings.create( + model="multimodalembedding@001", + input = "data:image/jpeg;base64,...", +) + +print(response) +``` + + + + + +Requests with GCS Image / Video URI +```python +from langchain_openai import OpenAIEmbeddings + +embeddings_models = "multimodalembedding@001" + +embeddings = OpenAIEmbeddings( + model="multimodalembedding@001", + base_url="http://0.0.0.0:4000", + api_key="sk-1234", # type: ignore +) + + +query_result = embeddings.embed_query( + "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png" +) +print(query_result) + +``` + +Requests with base64 encoded images + +```python +from langchain_openai import OpenAIEmbeddings + +embeddings_models = "multimodalembedding@001" + +embeddings = OpenAIEmbeddings( + model="multimodalembedding@001", + base_url="http://0.0.0.0:4000", + api_key="sk-1234", # type: ignore +) + + +query_result = embeddings.embed_query( + "data:image/jpeg;base64,..." +) +print(query_result) + +``` + + + + + + + + + +1. Add model to config.yaml +```yaml +default_vertex_config: + vertex_project: "adroit-crow-413218" + vertex_location: "us-central1" + vertex_credentials: adroit-crow-413218-a956eef1a2a8.json +``` + +2. Start Proxy + +``` +$ litellm --config /path/to/config.yaml +``` + +3. Make Request use OpenAI Python SDK + +```python +import vertexai + +from vertexai.vision_models import Image, MultiModalEmbeddingModel, Video +from vertexai.vision_models import VideoSegmentConfig +from google.auth.credentials import Credentials + + +LITELLM_PROXY_API_KEY = "sk-1234" +LITELLM_PROXY_BASE = "http://0.0.0.0:4000/vertex-ai" + +import datetime + +class CredentialsWrapper(Credentials): + def __init__(self, token=None): + super().__init__() + self.token = token + self.expiry = None # or set to a future date if needed + + def refresh(self, request): + pass + + def apply(self, headers, token=None): + headers['Authorization'] = f'Bearer {self.token}' + + @property + def expired(self): + return False # Always consider the token as non-expired + + @property + def valid(self): + return True # Always consider the credentials as valid + +credentials = CredentialsWrapper(token=LITELLM_PROXY_API_KEY) + +vertexai.init( + project="adroit-crow-413218", + location="us-central1", + api_endpoint=LITELLM_PROXY_BASE, + credentials = credentials, + api_transport="rest", + +) + +model = MultiModalEmbeddingModel.from_pretrained("multimodalembedding") +image = Image.load_from_file( + "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png" +) + +embeddings = model.get_embeddings( + image=image, + contextual_text="Colosseum", + dimension=1408, +) +print(f"Image Embedding: {embeddings.image_embedding}") +print(f"Text Embedding: {embeddings.text_embedding}") +``` + + + + + +### Text + Image + Video Embeddings + + + + +Text + Image + +```python +response = await litellm.aembedding( + model="vertex_ai/multimodalembedding@001", + input=["hey", "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"] # will be sent as a gcs image +) +``` + +Text + Video + +```python +response = await litellm.aembedding( + model="vertex_ai/multimodalembedding@001", + input=["hey", "gs://my-bucket/embeddings/supermarket-video.mp4"] # will be sent as a gcs image +) +``` + +Image + Video + +```python +response = await litellm.aembedding( + model="vertex_ai/multimodalembedding@001", + input=["gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png", "gs://my-bucket/embeddings/supermarket-video.mp4"] # will be sent as a gcs image +) +``` + + + + + +1. Add model to config.yaml +```yaml +model_list: + - model_name: multimodalembedding@001 + litellm_params: + model: vertex_ai/multimodalembedding@001 + vertex_project: "adroit-crow-413218" + vertex_location: "us-central1" + vertex_credentials: adroit-crow-413218-a956eef1a2a8.json + +litellm_settings: + drop_params: True +``` + +2. Start Proxy + +``` +$ litellm --config /path/to/config.yaml +``` + +3. Make Request use OpenAI Python SDK, Langchain Python SDK + + +Text + Image + +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +# # request sent to model set on litellm proxy, `litellm --model` +response = client.embeddings.create( + model="multimodalembedding@001", + input = ["hey", "gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png"], +) + +print(response) +``` + +Text + Video +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +# # request sent to model set on litellm proxy, `litellm --model` +response = client.embeddings.create( + model="multimodalembedding@001", + input = ["hey", "gs://my-bucket/embeddings/supermarket-video.mp4"], +) + +print(response) +``` + +Image + Video +```python +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +# # request sent to model set on litellm proxy, `litellm --model` +response = client.embeddings.create( + model="multimodalembedding@001", + input = ["gs://cloud-samples-data/vertex-ai/llm/prompts/landmark1.png", "gs://my-bucket/embeddings/supermarket-video.mp4"], +) + +print(response) +``` + + + + + ## **Gemini TTS (Text-to-Speech) Audio Output** :::info diff --git a/docs/my-website/docs/providers/vertex_embedding.md b/docs/my-website/docs/providers/vertex_embedding.md deleted file mode 100644 index 5656ade337..0000000000 --- a/docs/my-website/docs/providers/vertex_embedding.md +++ /dev/null @@ -1,587 +0,0 @@ -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 19d7faf796..99472981f0 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -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", diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index b40f0a72a5..7932881f48 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -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 = { diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 02b0f79280..2c53457736 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -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}" 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 dabc620a6d..70b068b5a4 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 @@ -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( diff --git a/litellm/llms/vertex_ai/vertex_embeddings/bge.py b/litellm/llms/vertex_ai/vertex_embeddings/bge.py deleted file mode 100644 index 2eff0ba96d..0000000000 --- a/litellm/llms/vertex_ai/vertex_embeddings/bge.py +++ /dev/null @@ -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 - diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index 5a3a4a7188..97af558041 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -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"] diff --git a/litellm/llms/vertex_ai/vertex_embeddings/types.py b/litellm/llms/vertex_ai/vertex_embeddings/types.py index fa9794d79a..7f85ea46f3 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/types.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/types.py @@ -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]] diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/main.py b/litellm/llms/vertex_ai/vertex_gemma_models/main.py index 41bd6b5431..8203b285eb 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, 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( diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index ce50bf311e..9ddbc461a7 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -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( diff --git a/litellm/llms/vertex_ai/vertex_model_garden/main.py b/litellm/llms/vertex_ai/vertex_model_garden/main.py index fe7d0862e0..1c57096734 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, 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( diff --git a/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py b/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py deleted file mode 100644 index 156ab95184..0000000000 --- a/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py +++ /dev/null @@ -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 - - 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 deleted file mode 100644 index 20150501ad..0000000000 --- a/tests/test_litellm/llms/vertex_ai/test_bge_response_transformation.py +++ /dev/null @@ -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 - ) - 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 deleted file mode 100644 index 46f365094c..0000000000 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_psc_endpoint_support.py +++ /dev/null @@ -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}" -