diff --git a/docs/my-website/docs/pass_through/vertex_ai.md b/docs/my-website/docs/pass_through/vertex_ai.md index 61f47aab6de..1bf55582301 100644 --- a/docs/my-website/docs/pass_through/vertex_ai.md +++ b/docs/my-website/docs/pass_through/vertex_ai.md @@ -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 + + + + +#### 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="", # enter your project id + 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") + +``` + + + + + + #### 1. Set `default_vertex_config` on your `config.yaml` @@ -95,12 +138,43 @@ response = model.generate_content( print(response.text) ``` + + + + ## Usage Examples ### Gemini API (Generate Content) - + + +```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) +``` + + + ```python import vertexai @@ -171,7 +245,45 @@ curl http://localhost:4000/vertex-ai/publishers/google/models/gemini-1.5-flash-0 ### Embeddings API - + + + +```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] +``` + + + + ```python from typing import List, Optional @@ -249,7 +361,54 @@ curl http://localhost:4000/vertex-ai/publishers/google/models/textembedding-geck ### Imagen API - + + + + +```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") + +``` + + + ```python from typing import List, Optional @@ -338,7 +497,50 @@ curl http://localhost:4000/vertex-ai/publishers/google/models/imagen-3.0-generat - + + + +```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}") +``` + + + + + ```python from typing import List, Optional @@ -425,7 +627,47 @@ Create Fine Tuning Job - + + +```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) + +``` + + + + ```python from typing import List, Optional diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index aff51624afe..7c78eb5865a 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -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 diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 433480bc4ef..00b89edb953 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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 diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index ac17a1d84c4..21beb965cc0 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -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: diff --git a/litellm/proxy/tests/test_vertex_sdk_forward_headers.py b/litellm/proxy/tests/test_vertex_sdk_forward_headers.py new file mode 100644 index 00000000000..b291be438ca --- /dev/null +++ b/litellm/proxy/tests/test_vertex_sdk_forward_headers.py @@ -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") diff --git a/litellm/proxy/vertex_ai_endpoints/vertex_endpoints.py b/litellm/proxy/vertex_ai_endpoints/vertex_endpoints.py index 53edbbcfd3c..fe1f46bda41 100644 --- a/litellm/proxy/vertex_ai_endpoints/vertex_endpoints.py +++ b/litellm/proxy/vertex_ai_endpoints/vertex_endpoints.py @@ -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)