phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 19:54:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add example using litellm proxy with gemini context caching

Browse source
This commit is contained in:
Ishaan Jaff 2024-08-08 11:35:41 -07:00 • committed by Krrish Dholakia
parent 7d37ff8070
commit 2be9c185e5
1 changed files with 20 additions and 36 deletions
Download patch file Download diff file Expand all files Collapse all files

56
litellm/proxy/tests/test_gemini_context_caching.py
View file

@ -1,54 +1,38 @@
import datetime import datetime
import httpx
import openai 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 here client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
LITELLM_BASE_URL = "http://0.0.0.0:4000" vertexai.init(project="adroit-crow-413218", location="us-central1")
LITELLM_PROXY_API_KEY = "sk-1234"
client = openai.OpenAI(api_key=LITELLM_PROXY_API_KEY, base_url=LITELLM_BASE_URL)
httpx_client = httpx.Client(timeout=30)
################################
# First create a cachedContents object
print("creating cached content") print("creating cached content")
create_cache = httpx_client.post( contents_here: list[Content] = [
url=f"{LITELLM_BASE_URL}/vertex-ai/cachedContents", Content(role="user", parts=[Part.from_text("huge string of text here" * 10000)])
headers={"Authorization": f"Bearer {LITELLM_PROXY_API_KEY}"}, ]
json={ cached_content = caching.CachedContent.create(
"model": "gemini-1.5-pro-001", model_name="gemini-1.5-pro-001",
"contents": [ contents=contents_here,
{ expire_time=datetime.datetime(2024, 8, 10),
"role": "user",
"parts": [
{
"text": "This is sample text to demonstrate explicit caching."
* 4000
}
],
}
],
},
) )
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"]
################################# created_Caches = caching.CachedContent.list()
# Use the `cachedContents` object in your /chat/completions
print("created_Caches contents=", created_Caches)
response = client.chat.completions.create( # type: ignore response = client.chat.completions.create( # type: ignore
model="gemini-1.5-pro-001", model="gemini-1.5-pro-001",
max_tokens=8192, max_tokens=8192,
messages=[ messages=[
{ {
"role": "user", "role": "user",
"content": "what is the sample text about?", "content": "quote all everything above this message",
}, },
], ],
temperature="0.7", temperature=0.7,
extra_body={"cached_content": cached_content_name}, # 👈 key change extra_body={"cached_content": cached_content.resource_name},
) )
print("response from proxy", response) print("response from proxy", response)