Skip to content

Commit 772a837

Browse files
Valeriy-Burlakariathakkar
authored andcommitted
feat(generative-ai): Add new sample for Context Cache's list feature (GoogleCloudPlatform#12596)
* Gen AI: Add sample for the 'get list of content caches' section * review comments: make sample simpler * review comments: rename the sample to follow standard naming practice * lint failure: rollback import statements re-arrangement * Update comments to follow the runnable-snippets-plan standard * refresh CI builds
1 parent a34f85b commit 772a837

File tree

2 files changed

+75
-0
lines changed

2 files changed

+75
-0
lines changed

generative_ai/context_caching/context_caching_test.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,11 @@
1616

1717
from typing import Generator
1818

19+
1920
import create_context_cache
2021
import delete_context_cache
2122
import get_context_cache
23+
import list_content_caches
2224
import pytest
2325
import update_context_cache
2426
import use_context_cache
@@ -48,6 +50,12 @@ def test_get_context_cache(cache_id: str) -> None:
4850
assert response
4951

5052

53+
def test_get_list_of_context_caches(cache_id: str) -> None:
54+
response = list_content_caches.list_content_caches()
55+
cache_id_is_in_response = any([cc for cc in response if cc.name == cache_id])
56+
assert cache_id_is_in_response
57+
58+
5159
def test_update_context_cache(cache_id: str) -> None:
5260
response = update_context_cache.update_context_cache(cache_id)
5361
assert response
Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
import os
15+
16+
from typing import List
17+
18+
from google.cloud.aiplatform_v1beta1.types.cached_content import CachedContent
19+
20+
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
21+
22+
23+
def list_content_caches() -> List[CachedContent]:
24+
# [START generativeaionvertexai_gemini_get_list_of_context_caches]
25+
import json
26+
27+
import vertexai
28+
29+
from vertexai.preview import caching
30+
31+
# TODO(developer): Update & uncomment line below
32+
# PROJECT_ID = "your-project-id"
33+
vertexai.init(project=PROJECT_ID, location="us-central1")
34+
35+
cached_content_list = caching.CachedContent.list()
36+
# Access individual properties of a CachedContent object
37+
for cc in cached_content_list:
38+
print(
39+
f"Cached content '{cc.display_name}' for model '{cc.model_name}' expires at {cc.expire_time}."
40+
)
41+
# Example response:
42+
# Cached content 'scientific-articles' for model '.../gemini-1.5-pro-001' expires at 2024-09-23 18:01:47.242036+00:00.
43+
44+
# or convert the ContentCache object to a dictionary:
45+
for cc in cached_content_list:
46+
print(json.dumps(cc.to_dict(), indent=2))
47+
# Example response:
48+
# {
49+
# "name": "projects/[PROJECT_NUMBER]/locations/us-central1/cachedContents/4101407070023057408",
50+
# "model": "projects/[PROJECT_ID]/locations/us-central1/publishers/google/models/gemini-1.5-pro-001",
51+
# "createTime": "2024-09-16T12:41:09.998635Z",
52+
# "updateTime": "2024-09-16T12:41:09.998635Z",
53+
# "expireTime": "2024-09-16T13:41:09.989729Z",
54+
# "displayName": "scientific-articles"
55+
# "usageMetadata": {
56+
# "totalTokenCount": 43130,
57+
# "textCount": 153,
58+
# "imageCount": 167
59+
# }
60+
# }
61+
62+
# [END generativeaionvertexai_gemini_get_list_of_context_caches]
63+
return cached_content_list
64+
65+
66+
if __name__ == "__main__":
67+
list_content_caches()

0 commit comments

Comments
 (0)