mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-12 21:04:10 +00:00
doc on using litellm proxy with vertex ai content caching
This commit is contained in:
committed by
Krrish Dholakia
parent
c723610c59
commit
bf38d14201
@@ -440,20 +440,12 @@ Use Vertex AI Context Caching
|
||||
1. Add model to config.yaml
|
||||
```yaml
|
||||
model_list:
|
||||
# used for /chat/completions, /completions, /embeddings endpoints
|
||||
- model_name: gemini-1.5-pro-001
|
||||
litellm_params:
|
||||
model: vertex_ai_beta/gemini-1.5-pro-001
|
||||
vertex_project: "project-id"
|
||||
vertex_location: "us-central1"
|
||||
vertex_credentials: "adroit-crow-413218-a956eef1a2a8.json" # Add path to service account.json
|
||||
|
||||
# used for the /cachedContent and vertexAI native endpoints
|
||||
default_vertex_config:
|
||||
vertex_project: "adroit-crow-413218"
|
||||
vertex_location: "us-central1"
|
||||
vertex_credentials: "adroit-crow-413218-a956eef1a2a8.json" # Add path to service account.json
|
||||
|
||||
vertex_credentials: "/path/to/service_account.json" # [OPTIONAL] Do this OR `!gcloud auth application-default login` - run this to add vertex credentials to your env
|
||||
```
|
||||
|
||||
2. Start Proxy
|
||||
@@ -463,71 +455,45 @@ $ litellm --config /path/to/config.yaml
|
||||
```
|
||||
|
||||
3. Make Request!
|
||||
We make the request in two steps:
|
||||
- Create a cachedContents object
|
||||
- Use the cachedContents object in your /chat/completions
|
||||
|
||||
**Create a cachedContents object**
|
||||
|
||||
First, create a cachedContents object by calling the Vertex `cachedContents` endpoint. The LiteLLM proxy forwards the `/cachedContents` request to the VertexAI API.
|
||||
|
||||
```python
|
||||
import httpx
|
||||
import datetime
|
||||
import openai
|
||||
import vertexai
|
||||
from vertexai.generative_models import Content, Part
|
||||
from vertexai.preview import caching
|
||||
from vertexai.preview.generative_models import GenerativeModel
|
||||
|
||||
# Set Litellm proxy variables
|
||||
LITELLM_BASE_URL = "http://0.0.0.0:4000"
|
||||
LITELLM_PROXY_API_KEY = "sk-1234"
|
||||
|
||||
httpx_client = httpx.Client(timeout=30)
|
||||
|
||||
print("Creating cached content")
|
||||
create_cache = httpx_client.post(
|
||||
url=f"{LITELLM_BASE_URL}/vertex-ai/cachedContents",
|
||||
headers={"Authorization": f"Bearer {LITELLM_PROXY_API_KEY}"},
|
||||
json={
|
||||
"model": "gemini-1.5-pro-001",
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [{
|
||||
"text": "This is sample text to demonstrate explicit caching." * 4000
|
||||
}]
|
||||
}
|
||||
],
|
||||
}
|
||||
# use Vertex AI SDK to create CachedContent
|
||||
vertexai.init(project="adroit-crow-413218", location="us-central1")
|
||||
print("creating cached content")
|
||||
contents_here: list[Content] = [
|
||||
Content(role="user", parts=[Part.from_text("huge string of text here" * 10000)])
|
||||
]
|
||||
cached_content = caching.CachedContent.create(
|
||||
model_name="gemini-1.5-pro-001",
|
||||
contents=contents_here,
|
||||
expire_time=datetime.datetime(2024, 8, 10),
|
||||
)
|
||||
|
||||
print("Response from create_cache:", create_cache)
|
||||
create_cache_response = create_cache.json()
|
||||
print("JSON from create_cache:", create_cache_response)
|
||||
cached_content_name = create_cache_response["name"]
|
||||
```
|
||||
|
||||
**Use the cachedContents object in your /chat/completions request to VertexAI**
|
||||
|
||||
```python
|
||||
import openai
|
||||
|
||||
# Set Litellm proxy variables
|
||||
LITELLM_BASE_URL = "http://0.0.0.0:4000"
|
||||
LITELLM_PROXY_API_KEY = "sk-1234"
|
||||
|
||||
client = openai.OpenAI(api_key=LITELLM_PROXY_API_KEY, base_url=LITELLM_BASE_URL)
|
||||
|
||||
# use OpenAI SDK to send a request to LiteLLM Proxy
|
||||
# base_url is litellm proxy server and api_key is api key to litellm proxy
|
||||
client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
|
||||
response = client.chat.completions.create(
|
||||
model="gemini-1.5-pro-001",
|
||||
max_tokens=8192,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What is the sample text about?",
|
||||
"content": "hello!",
|
||||
},
|
||||
],
|
||||
temperature=0.7,
|
||||
extra_body={"cached_content": cached_content_name}, # Use the cached content
|
||||
temperature="0.7",
|
||||
extra_body={"cached_content": cached_content.resource_name},
|
||||
)
|
||||
|
||||
print("Response from proxy:", response)
|
||||
print("response from proxy", response)
|
||||
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
Reference in New Issue
Block a user