Skip to content

Commit e1d8c7a

Browse files
Validate name checks for CachedContent creation
Change-Id: Ie41602621d99ddff6404c6708c7278e0da790652
1 parent 2cde1a2 commit e1d8c7a

File tree

2 files changed

+14
-1
lines changed

2 files changed

+14
-1
lines changed

google/generativeai/caching.py

+3-1
Original file line numberDiff line numberDiff line change
@@ -88,8 +88,10 @@ def _prepare_create_request(
8888
ttl: Optional[caching_types.ExpirationTypes] = datetime.timedelta(hours=1),
8989
) -> glm.CreateCachedContentRequest:
9090
"""Prepares a CreateCachedContentRequest."""
91-
if "cachedContents/" not in name:
91+
if caching_types.valid_cached_content_name(name):
9292
name = "cachedContents/" + name
93+
else:
94+
raise ValueError(caching_types.NAME_ERROR_MESSAGE.format(name=name))
9395

9496
if "/" not in model:
9597
model = "models/" + model

google/generativeai/types/caching_types.py

+11
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,21 @@
1717
import datetime
1818
from typing import Optional, Union
1919
from typing_extensions import TypedDict
20+
import re
2021

2122
__all__ = ["TTL"]
2223

2324

25+
_VALID_CACHED_CONTENT_NAME = r"([a-z0-9-\.]+)$"
26+
NAME_ERROR_MESSAGE = (
27+
"The `name` must consist of alphanumeric characters (or `-` or `.`). Received: `{name}`"
28+
)
29+
30+
31+
def valid_cached_content_name(name: str) -> bool:
32+
return re.match(_VALID_CACHED_CONTENT_NAME, name) is not None
33+
34+
2435
class TTL(TypedDict):
2536
seconds: int
2637

0 commit comments

Comments
 (0)