From bd11434135310c059d3fc1e42c091b8265fd9de7 Mon Sep 17 00:00:00 2001 From: mukunda katta Date: Sun, 19 Apr 2026 17:26:45 -0700 Subject: [PATCH] feat(caches): support labels in cached content config --- google/genai/caches.py | 3 ++ .../genai/tests/caches/test_cache_labels.py | 43 +++++++++++++++++++ google/genai/types.py | 14 ++++++ 3 files changed, 60 insertions(+) create mode 100644 google/genai/tests/caches/test_cache_labels.py diff --git a/google/genai/caches.py b/google/genai/caches.py index 78ad51921..ad1e1053d 100644 --- a/google/genai/caches.py +++ b/google/genai/caches.py @@ -197,6 +197,9 @@ def _CreateCachedContentConfig_to_vertex( if getv(from_object, ['display_name']) is not None: setv(parent_object, ['displayName'], getv(from_object, ['display_name'])) + if getv(from_object, ['labels']) is not None: + setv(parent_object, ['labels'], getv(from_object, ['labels'])) + if getv(from_object, ['contents']) is not None: setv( parent_object, diff --git a/google/genai/tests/caches/test_cache_labels.py b/google/genai/tests/caches/test_cache_labels.py new file mode 100644 index 000000000..6c4b882f9 --- /dev/null +++ b/google/genai/tests/caches/test_cache_labels.py @@ -0,0 +1,43 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from ... import caches +from ... import types + + +def test_create_cached_content_config_to_vertex_includes_labels(): + parent_object = {} + + caches._CreateCachedContentConfig_to_vertex( + types.CreateCachedContentConfig( + display_name='test cache', + labels={'team': 'genai', 'use_case': 'billing'}, + ), + parent_object, + ) + + assert parent_object['labels'] == { + 'team': 'genai', + 'use_case': 'billing', + } + + +def test_cached_content_accepts_labels(): + cached_content = types.CachedContent( + name='cachedContents/123', + labels={'team': 'genai'}, + ) + + assert cached_content.labels == {'team': 'genai'} diff --git a/google/genai/types.py b/google/genai/types.py index 20c4f1c0e..8fe14db08 100644 --- a/google/genai/types.py +++ b/google/genai/types.py @@ -14333,6 +14333,10 @@ class CreateCachedContentConfig(_common.BaseModel): description="""The user-generated meaningful display name of the cached content. """, ) + labels: Optional[dict[str, str]] = Field( + default=None, + description="""User specified labels to track billing usage.""", + ) contents: Optional[ContentListUnion] = Field( default=None, description="""The content to cache. @@ -14383,6 +14387,9 @@ class CreateCachedContentConfigDict(TypedDict, total=False): """The user-generated meaningful display name of the cached content. """ + labels: Optional[dict[str, str]] + """User specified labels to track billing usage.""" + contents: Optional[ContentListUnionDict] """The content to cache. """ @@ -14506,6 +14513,10 @@ class CachedContent(_common.BaseModel): default=None, description="""The user-generated meaningful display name of the cached content.""", ) + labels: Optional[dict[str, str]] = Field( + default=None, + description="""User specified labels to track billing usage.""", + ) model: Optional[str] = Field( default=None, description="""The name of the publisher model to use for cached content.""", @@ -14535,6 +14546,9 @@ class CachedContentDict(TypedDict, total=False): display_name: Optional[str] """The user-generated meaningful display name of the cached content.""" + labels: Optional[dict[str, str]] + """User specified labels to track billing usage.""" + model: Optional[str] """The name of the publisher model to use for cached content."""