From 9078f075f9b1322a5cb3880aa9890ef6cad181fd Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 21 Aug 2024 17:35:31 -0700 Subject: [PATCH] docs add example using litellm with vertex python sdk --- .../my-website/docs/pass_through/vertex_ai.md | 191 +++++++++++++++++- 1 file changed, 186 insertions(+), 5 deletions(-) diff --git a/docs/my-website/docs/pass_through/vertex_ai.md b/docs/my-website/docs/pass_through/vertex_ai.md index 53ae410913e..bfc1bc98b84 100644 --- a/docs/my-website/docs/pass_through/vertex_ai.md +++ b/docs/my-website/docs/pass_through/vertex_ai.md @@ -1,3 +1,7 @@ +import Image from '@theme/IdealImage'; +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + # [BETA] Vertex AI Endpoints (Pass-Through) Pass-through endpoints for Vertex AI - call provider-specific endpoint, in native format (no translation). @@ -40,16 +44,120 @@ litellm --config /path/to/config.yaml #### 3. Test it -```shell -curl http://localhost:4000/vertex-ai/publishers/google/models/textembedding-gecko@001:countTokens \ --H "Content-Type: application/json" \ --H "Authorization: Bearer sk-1234" \ --d '{"instances":[{"content": "gm"}]}' +```python +import vertexai +from google.auth.credentials import Credentials +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 + + +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", + request_metadata=[("Authorization", f"Bearer {LITELLM_PROXY_API_KEY}")], +) + +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) ``` + ## Usage Examples ### Gemini API (Generate Content) + + + +```python +import vertexai +from google.auth.credentials import Credentials +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 + + +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", + request_metadata=[("Authorization", f"Bearer {LITELLM_PROXY_API_KEY}")], +) + +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) +``` + + + + ```shell curl http://localhost:4000/vertex-ai/publishers/google/models/gemini-1.5-flash-001:generateContent \ -H "Content-Type: application/json" \ @@ -57,8 +165,77 @@ curl http://localhost:4000/vertex-ai/publishers/google/models/gemini-1.5-flash-0 -d '{"contents":[{"role": "user", "parts":[{"text": "hi"}]}]}' ``` + + + + ### Embeddings API + + + +```python +from typing import List, Optional +from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel +import vertexai +from google.auth.credentials import Credentials +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 + + +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", + request_metadata=[("Authorization", f"Bearer {LITELLM_PROXY_API_KEY}")], + + +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] +``` + + + + ```shell curl http://localhost:4000/vertex-ai/publishers/google/models/textembedding-gecko@001:predict \ -H "Content-Type: application/json" \ @@ -66,6 +243,10 @@ curl http://localhost:4000/vertex-ai/publishers/google/models/textembedding-geck -d '{"instances":[{"content": "gm"}]}' ``` + + + + ### Imagen API ```shell