scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

Merge pull request #5435 from BerriAI/litellm_fwd_vtx_sdk_headers

Browse source 
[Feat-Proxy] Pass through Vertex Endpoint - allow forwarding vertex credentials
This commit is contained in:
Ishaan Jaff 2024-08-29 17:24:35 -07:00 • committed by GitHub
commit ef16738720
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 326 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

252
docs/my-website/docs/pass_through/vertex_ai.md
View file

@ -22,8 +22,51 @@ Looking for the Unified API (OpenAI format) for VertexAI ? [Go here - using vert
- Tuning API
- CountTokens API
## Authentication to Vertex AI
LiteLLM Proxy Server supports two methods of authentication to Vertex AI:
1. Pass Vertex Credetials client side to proxy server
2. Set Vertex AI credentials on proxy server
## Quick Start Usage
<Tabs>
<TabItem value="without_default_config" label="Pass Vertex Credetials client side to proxy server">
#### 1. Start litellm proxy
```shell
litellm --config /path/to/config.yaml
```
#### 2. Test it
```python
import vertexai
from vertexai.preview.generative_models import GenerativeModel
LITE_LLM_ENDPOINT = "http://localhost:4000"
vertexai.init(
project="<your-vertex-ai-project-id>", # enter your project id
location="<your-vertex-ai-location>", # enter your region
api_endpoint=f"{LITE_LLM_ENDPOINT}/vertex-ai", # route on litellm
api_transport="rest",
)
model = GenerativeModel(model_name="gemini-1.0-pro")
model.generate_content("hi")
```
</TabItem>
<TabItem value="with_default_config" label="Set Vertex AI Credentials on Proxy Server">
#### 1. Set `default_vertex_config` on your `config.yaml`
@ -95,12 +138,43 @@ response = model.generate_content(
print(response.text)
```
</TabItem>
</Tabs>
## Usage Examples
### Gemini API (Generate Content)
<Tabs>
<TabItem value="py" label="Vertex Python SDK">
<TabItem value="client_side" label="Vertex Python SDK (client side vertex credentials)">
```python
import vertexai
from vertexai.generative_models import GenerativeModel
LITELLM_PROXY_API_KEY = "sk-1234"
LITELLM_PROXY_BASE = "http://0.0.0.0:4000/vertex-ai"
vertexai.init(
project="adroit-crow-413218",
location="us-central1",
api_endpoint=LITELLM_PROXY_BASE,
api_transport="rest",
)
model = GenerativeModel("gemini-1.5-flash-001")
response = model.generate_content(
"What's a good name for a flower shop that specializes in selling bouquets of dried flowers?"
)
print(response.text)
```
</TabItem>
<TabItem value="py" label="Vertex Python SDK (litellm virtual keys client side)">
```python
import vertexai
@ -171,7 +245,45 @@ curl http://localhost:4000/vertex-ai/publishers/google/models/gemini-1.5-flash-0
### Embeddings API
<Tabs>
<TabItem value="py" label="Vertex Python SDK">
<TabItem value="client_side" label="Vertex Python SDK (client side vertex credentials)">
```python
from typing import List, Optional
from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
import vertexai
from vertexai.generative_models import GenerativeModel
LITELLM_PROXY_API_KEY = "sk-1234"
LITELLM_PROXY_BASE = "http://0.0.0.0:4000/vertex-ai"
import datetime
vertexai.init(
project="adroit-crow-413218",
location="us-central1",
api_endpoint=LITELLM_PROXY_BASE,
api_transport="rest",
)
def embed_text(
texts: List[str] = ["banana muffins? ", "banana bread? banana muffins?"],
task: str = "RETRIEVAL_DOCUMENT",
model_name: str = "text-embedding-004",
dimensionality: Optional[int] = 256,
) -> List[List[float]]:
"""Embeds texts with a pre-trained, foundational model."""
model = TextEmbeddingModel.from_pretrained(model_name)
inputs = [TextEmbeddingInput(text, task) for text in texts]
kwargs = dict(output_dimensionality=dimensionality) if dimensionality else {}
embeddings = model.get_embeddings(inputs, **kwargs)
return [embedding.values for embedding in embeddings]
```
</TabItem>
<TabItem value="py" label="Vertex Python SDK (litellm virtual keys client side)">
```python
from typing import List, Optional
@ -249,7 +361,54 @@ curl http://localhost:4000/vertex-ai/publishers/google/models/textembedding-geck
### Imagen API
<Tabs>
<TabItem value="py" label="Vertex Python SDK">
<TabItem value="client_side" label="Vertex Python SDK (client side vertex credentials)">
```python
from typing import List, Optional
from vertexai.preview.vision_models import ImageGenerationModel
import vertexai
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
vertexai.init(
project="adroit-crow-413218",
location="us-central1",
api_endpoint=LITELLM_PROXY_BASE,
api_transport="rest",
)
model = ImageGenerationModel.from_pretrained("imagen-3.0-generate-001")
images = model.generate_images(
prompt=prompt,
# Optional parameters
number_of_images=1,
language="en",
# You can't use a seed value and watermark at the same time.
# add_watermark=False,
# seed=100,
aspect_ratio="1:1",
safety_filter_level="block_some",
person_generation="allow_adult",
)
images[0].save(location=output_file, include_generation_parameters=False)
# Optional. View the generated image in a notebook.
# images[0].show()
print(f"Created output image using {len(images[0]._image_bytes)} bytes")
```
</TabItem>
<TabItem value="py" label="Vertex Python SDK (litellm virtual keys client side)">
```python
from typing import List, Optional
@ -338,7 +497,50 @@ curl http://localhost:4000/vertex-ai/publishers/google/models/imagen-3.0-generat
<Tabs>
<TabItem value="py" label="Vertex Python SDK">
<TabItem value="client_side" label="Vertex Python SDK (client side vertex credentials)">
```python
from typing import List, Optional
from vertexai.generative_models import GenerativeModel
import vertexai
LITELLM_PROXY_API_KEY = "sk-1234"
LITELLM_PROXY_BASE = "http://0.0.0.0:4000/vertex-ai"
import datetime
vertexai.init(
project="adroit-crow-413218",
location="us-central1",
api_endpoint=LITELLM_PROXY_BASE,
api_transport="rest",
)
model = GenerativeModel("gemini-1.5-flash-001")
prompt = "Why is the sky blue?"
# Prompt tokens count
response = model.count_tokens(prompt)
print(f"Prompt Token Count: {response.total_tokens}")
print(f"Prompt Character Count: {response.total_billable_characters}")
# Send text to Gemini
response = model.generate_content(prompt)
# Response tokens count
usage_metadata = response.usage_metadata
print(f"Prompt Token Count: {usage_metadata.prompt_token_count}")
print(f"Candidates Token Count: {usage_metadata.candidates_token_count}")
print(f"Total Token Count: {usage_metadata.total_token_count}")
```
</TabItem>
<TabItem value="py" label="Vertex Python SDK (litellm virtual keys client side)">
```python
from typing import List, Optional
@ -425,7 +627,47 @@ Create Fine Tuning Job
<Tabs>
<TabItem value="py" label="Vertex Python SDK">
<TabItem value="client_side" label="Vertex Python SDK (client side vertex credentials)">
```python
from typing import List, Optional
from vertexai.preview.tuning import sft
import vertexai
LITELLM_PROXY_API_KEY = "sk-1234"
LITELLM_PROXY_BASE = "http://0.0.0.0:4000/vertex-ai"
vertexai.init(
project="adroit-crow-413218",
location="us-central1",
api_endpoint=LITELLM_PROXY_BASE,
api_transport="rest",
)
# TODO(developer): Update project
vertexai.init(project=PROJECT_ID, location="us-central1")
sft_tuning_job = sft.train(
source_model="gemini-1.0-pro-002",
train_dataset="gs://cloud-samples-data/ai-platform/generative_ai/sft_train_data.jsonl",
)
# Polling for job completion
while not sft_tuning_job.has_ended:
time.sleep(60)
sft_tuning_job.refresh()
print(sft_tuning_job.tuned_model_name)
print(sft_tuning_job.tuned_model_endpoint_name)
print(sft_tuning_job.experiment)
```
</TabItem>
<TabItem value="py" label="Vertex Python SDK (litellm virtual keys client side)">
```python
from typing import List, Optional

21
litellm/proxy/auth/auth_utils.py
View file

@ -227,3 +227,24 @@ def get_key_model_tpm_limit(user_api_key_dict: UserAPIKeyAuth) -> Optional[dict]
return user_api_key_dict.metadata["model_tpm_limit"]
return None
def is_pass_through_provider_route(route: str) -> bool:
PROVIDER_SPECIFIC_PASS_THROUGH_ROUTES = [
"vertex-ai",
]
# check if any of the prefixes are in the route
for prefix in PROVIDER_SPECIFIC_PASS_THROUGH_ROUTES:
if prefix in route:
return True
return False
def should_run_auth_on_pass_through_provider_route(route: str) -> bool:
"""
Use this to decide if the rest of the LiteLLM Virtual Key auth checks should run on /vertex-ai/{endpoint} routes
"""
# by default we do not run virtual key auth checks on /vertex-ai/{endpoint} routes
return False

8
litellm/proxy/auth/user_api_key_auth.py
View file

@ -61,7 +61,9 @@ from litellm.proxy.auth.auth_utils import (
check_if_request_size_is_safe,
get_request_route,
is_llm_api_route,
is_pass_through_provider_route,
route_in_additonal_public_routes,
should_run_auth_on_pass_through_provider_route,
)
from litellm.proxy.auth.oauth2_check import check_oauth2_token
from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request
@ -204,7 +206,11 @@ async def user_api_key_auth(
):
# check if public endpoint
return UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY)
elif is_pass_through_provider_route(route=route):
if should_run_auth_on_pass_through_provider_route(route=route) is False:
return UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
)
if general_settings.get("enable_oauth2_auth", False) is True:
# return UserAPIKeyAuth object
# helper to check if the api_key is a valid oauth2 token

8
litellm/proxy/proxy_config.yaml
View file

@ -13,14 +13,6 @@ model_list:
model: cohere/rerank-english-v3.0
api_key: os.environ/COHERE_API_KEY
general_settings:
enable_oauth2_proxy_auth: True
oauth2_config_mappings:
token: X-Auth-Token
user_id: X-Auth-Client-ID
team_id: X-Auth-Team-ID
max_budget: X-Auth-Max-Budget
models: X-Auth-Allowed-Models
# default off mode
litellm_settings:

14
litellm/proxy/tests/test_vertex_sdk_forward_headers.py Normal file
View file

@ -0,0 +1,14 @@
import vertexai
from vertexai.preview.generative_models import GenerativeModel
LITE_LLM_ENDPOINT = "http://localhost:4000"
vertexai.init(
project="adroit-crow-413218",
location="us-central1",
api_endpoint=f"{LITE_LLM_ENDPOINT}/vertex-ai",
api_transport="rest",
)
model = GenerativeModel(model_name="gemini-1.0-pro")
model.generate_content("hi")

56
litellm/proxy/vertex_ai_endpoints/vertex_endpoints.py
View file

@ -84,31 +84,49 @@ async def vertex_proxy_route(
):
encoded_endpoint = httpx.URL(endpoint).path
import re
from litellm.fine_tuning.main import vertex_fine_tuning_apis_instance
verbose_proxy_logger.debug("requested endpoint %s", endpoint)
headers: dict = {}
# Use headers from the incoming request if default_vertex_config is not set
if default_vertex_config is None:
raise ValueError(
"Vertex credentials not added on litellm proxy, please add `default_vertex_config` on your config.yaml"
headers = dict(request.headers) or {}
verbose_proxy_logger.debug(
"default_vertex_config not set, incoming request headers %s", headers
)
vertex_project = default_vertex_config.get("vertex_project", None)
vertex_location = default_vertex_config.get("vertex_location", None)
vertex_credentials = default_vertex_config.get("vertex_credentials", None)
base_target_url = f"https://{vertex_location}-aiplatform.googleapis.com/"
# extract location from endpoint, endpoint
# "v1beta1/projects/adroit-crow-413218/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent"
match = re.search(r"/locations/([^/]+)", endpoint)
vertex_location = match.group(1) if match else None
base_target_url = f"https://{vertex_location}-aiplatform.googleapis.com/"
headers.pop("content-length", None)
_new_headers = {
"Authorization": headers.get("authorization"),
}
headers = _new_headers
else:
vertex_project = default_vertex_config.get("vertex_project")
vertex_location = default_vertex_config.get("vertex_location")
vertex_credentials = default_vertex_config.get("vertex_credentials")
auth_header, _ = vertex_fine_tuning_apis_instance._get_token_and_url(
model="",
gemini_api_key=None,
vertex_credentials=vertex_credentials,
vertex_project=vertex_project,
vertex_location=vertex_location,
stream=False,
custom_llm_provider="vertex_ai_beta",
api_base="",
)
base_target_url = f"https://{vertex_location}-aiplatform.googleapis.com/"
headers = {
"Authorization": f"Bearer {auth_header}",
}
auth_header, _ = vertex_fine_tuning_apis_instance._get_token_and_url(
model="",
gemini_api_key=None,
vertex_credentials=vertex_credentials,
vertex_project=vertex_project,
vertex_location=vertex_location,
stream=False,
custom_llm_provider="vertex_ai_beta",
api_base="",
)
headers = {
"Authorization": f"Bearer {auth_header}",
}
request_route = encoded_endpoint
verbose_proxy_logger.debug("request_route %s", request_route)