From 0ea1f367d7302d410eee07408c12f03f712b7710 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 21 Aug 2024 17:46:23 -0700 Subject: [PATCH] use litellm proxy with vertex ai sdk --- .../my-website/docs/pass_through/vertex_ai.md | 236 +++++++++++++++++- docs/my-website/docs/providers/vertex.md | 2 +- litellm/proxy/tests/test_vtx_sdk_embedding.py | 1 - 3 files changed, 233 insertions(+), 6 deletions(-) diff --git a/docs/my-website/docs/pass_through/vertex_ai.md b/docs/my-website/docs/pass_through/vertex_ai.md index bfc1bc98b8..00ab85f2d3 100644 --- a/docs/my-website/docs/pass_through/vertex_ai.md +++ b/docs/my-website/docs/pass_through/vertex_ai.md @@ -84,7 +84,6 @@ vertexai.init( api_endpoint=LITELLM_PROXY_BASE, credentials=credentials, api_transport="rest", - request_metadata=[("Authorization", f"Bearer {LITELLM_PROXY_API_KEY}")], ) model = GenerativeModel("gemini-1.5-flash-001") @@ -143,7 +142,7 @@ vertexai.init( api_endpoint=LITELLM_PROXY_BASE, credentials=credentials, api_transport="rest", - request_metadata=[("Authorization", f"Bearer {LITELLM_PROXY_API_KEY}")], + ) model = GenerativeModel("gemini-1.5-flash-001") @@ -216,7 +215,7 @@ vertexai.init( api_endpoint=LITELLM_PROXY_BASE, credentials=credentials, api_transport="rest", - request_metadata=[("Authorization", f"Bearer {LITELLM_PROXY_API_KEY}")], +) def embed_text( @@ -249,6 +248,80 @@ 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 + + +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 = 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") + +``` + + + + + ```shell curl http://localhost:4000/vertex-ai/publishers/google/models/imagen-3.0-generate-001:predict \ -H "Content-Type: application/json" \ @@ -256,8 +329,86 @@ curl http://localhost:4000/vertex-ai/publishers/google/models/imagen-3.0-generat -d '{"instances":[{"prompt": "make an otter"}], "parameters": {"sampleCount": 1}}' ``` + + + + ### Count Tokens API + + + + + +```python +from typing import List, Optional +from vertexai.generative_models import GenerativeModel +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 + + +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 = 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}") +``` + + + + + + + ```shell curl http://localhost:4000/vertex-ai/publishers/google/models/gemini-1.5-flash-001:countTokens \ -H "Content-Type: application/json" \ @@ -265,10 +416,83 @@ curl http://localhost:4000/vertex-ai/publishers/google/models/gemini-1.5-flash-0 -d '{"contents":[{"role": "user", "parts":[{"text": "hi"}]}]}' ``` + + + ### Tuning API Create Fine Tuning Job + + + + +```python +from typing import List, Optional +from vertexai.preview.tuning import sft +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 + + +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", +) + + +# 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) +``` + + + + + ```shell curl http://localhost:4000/vertex-ai/tuningJobs \ -H "Content-Type: application/json" \ @@ -279,4 +503,8 @@ curl http://localhost:4000/vertex-ai/tuningJobs \ "training_dataset_uri": "gs://cloud-samples-data/ai-platform/generative_ai/sft_train_data.jsonl" } }' -``` \ No newline at end of file +``` + + + + diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index b8dad6ccdc..31dfb8ecda 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -1694,7 +1694,7 @@ vertexai.init( api_endpoint=LITELLM_PROXY_BASE, credentials = credentials, api_transport="rest", - request_metadata=[("Authorization", f"Bearer {LITELLM_PROXY_API_KEY}")], + ) model = MultiModalEmbeddingModel.from_pretrained("multimodalembedding") diff --git a/litellm/proxy/tests/test_vtx_sdk_embedding.py b/litellm/proxy/tests/test_vtx_sdk_embedding.py index a6468884f9..285743fca6 100644 --- a/litellm/proxy/tests/test_vtx_sdk_embedding.py +++ b/litellm/proxy/tests/test_vtx_sdk_embedding.py @@ -42,7 +42,6 @@ vertexai.init( api_endpoint=LITELLM_PROXY_BASE, credentials=credentials, api_transport="rest", - request_metadata=[("Authorization", f"Bearer {LITELLM_PROXY_API_KEY}")], ) model = MultiModalEmbeddingModel.from_pretrained("multimodalembedding")