mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
fix(vertex_ai): reuse known context caches instead of listing on every request (#46857)
Vertex AI and Gemini context caching now remember each cache they create or find until 10s before its expireTime, so a repeat cache-marked request goes straight to generateContent instead of listing every cachedContents first. Concurrent async requests for one prefix share one lookup or create When generation rejects a remembered cache with a 400, 403 or 404 that names the cached content, the request rebuilds the cache and sends once more. Any other 400, 403 or 404 drops the record without a retry Fixes #36395
This commit is contained in:
Ondra Zahradnik
GitHub
parent
808333682e
commit
06e3e32c88
6 changed files with 1173 additions and 160 deletions
156
litellm/llms/vertex_ai/context_caching/known_cached_contents.py
Normal file
156
litellm/llms/vertex_ai/context_caching/known_cached_contents.py
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
import asyncio
|
||||
import contextlib
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final, Literal
|
||||
|
||||
KNOWN_CACHE_EXPIRY_MARGIN_SECONDS: Final = 10.0
|
||||
MAX_KNOWN_CACHED_CONTENTS: Final = 10_000
|
||||
_CACHE_REJECTED_STATUS_CODES: Final = frozenset({400, 403, 404})
|
||||
_NAMES_CACHED_CONTENT: Final = re.compile(r"cached?\s*content", re.IGNORECASE)
|
||||
_FORBIDS_FIELDS_BESIDE_A_CACHE: Final = re.compile(r"system[_\s]instruction", re.IGNORECASE)
|
||||
|
||||
_RFC3339_FRACTION: Final = re.compile(r"\.(\d{1,6})\d*")
|
||||
|
||||
|
||||
def names_rejected_cache(status_code: int, message: str) -> bool:
|
||||
return (
|
||||
status_code in _CACHE_REJECTED_STATUS_CODES
|
||||
and _NAMES_CACHED_CONTENT.search(message) is not None
|
||||
and _FORBIDS_FIELDS_BESIDE_A_CACHE.search(message) is None
|
||||
)
|
||||
|
||||
|
||||
def parse_expire_time(expire_time: object) -> float | None:
|
||||
"""Parse a Vertex RFC 3339 ``expireTime`` (nanosecond precision, ``Z`` suffix) to epoch seconds."""
|
||||
if not isinstance(expire_time, str) or not expire_time:
|
||||
return None
|
||||
normalized: Final = _RFC3339_FRACTION.sub(lambda m: "." + m.group(1), expire_time.replace("Z", "+00:00"), count=1)
|
||||
try:
|
||||
parsed: Final = datetime.fromisoformat(normalized)
|
||||
except ValueError:
|
||||
return None
|
||||
if parsed.tzinfo is None:
|
||||
return parsed.replace(tzinfo=timezone.utc).timestamp()
|
||||
return parsed.timestamp()
|
||||
|
||||
|
||||
class KnownCachedContents:
|
||||
"""Process-local record of the cachedContents this process created or found.
|
||||
|
||||
Vertex meters every cachedContents list page against a per-minute, per-location quota,
|
||||
so listing on every cache-marked request exhausts it under load. A cache is remembered
|
||||
until ``KNOWN_CACHE_EXPIRY_MARGIN_SECONDS`` before the ``expireTime`` Google reported,
|
||||
and after that requests list again, as they did before the record existed.
|
||||
"""
|
||||
|
||||
def __init__(self, clock: Callable[[], float] = time.time) -> None:
|
||||
self._clock: Final = clock
|
||||
self._entries: dict[tuple[tuple[str, str], str], tuple[str, float]] = {} # mutable-ok: guarded by _lock
|
||||
self._lock: Final = threading.Lock()
|
||||
self._inflight: weakref.WeakKeyDictionary[ # mutable-ok: private, guarded by _lock
|
||||
asyncio.AbstractEventLoop, dict[tuple[tuple[str, str], str], asyncio.Future[str | None]]
|
||||
] = weakref.WeakKeyDictionary()
|
||||
|
||||
@staticmethod
|
||||
def scope(
|
||||
base_url: str,
|
||||
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
|
||||
headers: Mapping[str, str],
|
||||
vertex_project: str | None,
|
||||
vertex_location: str | None,
|
||||
) -> tuple[str, str]:
|
||||
"""Where a cache name is valid: a Vertex project and location, or the project a Gemini key belongs to.
|
||||
|
||||
A custom ``api_base`` can replace the whole Vertex path, so the project and location are kept beside the URL.
|
||||
The Gemini key is read from the headers actually sent, so an ``extra_headers`` override scopes the
|
||||
cache to the caller's project. It is kept as-is in this in-memory scope and never logged or persisted.
|
||||
"""
|
||||
if custom_llm_provider != "gemini":
|
||||
return (base_url, f"{vertex_project}/{vertex_location}")
|
||||
sent_keys: Final = [value for name, value in headers.items() if name.lower() == "x-goog-api-key"]
|
||||
return (base_url, sent_keys[-1] if sent_keys else "")
|
||||
|
||||
def has_not_expired(self, expire_time: object) -> bool:
|
||||
expires_at: Final = parse_expire_time(expire_time)
|
||||
return expires_at is None or expires_at > self._clock()
|
||||
|
||||
def get(self, scope: tuple[str, str], cache_key: str) -> str | None:
|
||||
with self._lock:
|
||||
entry: Final = self._entries.get((scope, cache_key))
|
||||
if entry is None:
|
||||
return None
|
||||
name, reuse_until = entry
|
||||
if reuse_until > self._clock():
|
||||
return name
|
||||
del self._entries[(scope, cache_key)]
|
||||
return None
|
||||
|
||||
def remember(self, scope: tuple[str, str], cache_key: str, name: str | None, expire_time: object) -> None:
|
||||
expires_at: Final = parse_expire_time(expire_time)
|
||||
if not name or expires_at is None:
|
||||
return
|
||||
reuse_until: Final = expires_at - KNOWN_CACHE_EXPIRY_MARGIN_SECONDS
|
||||
now: Final = self._clock()
|
||||
if reuse_until <= now:
|
||||
return
|
||||
with self._lock:
|
||||
if len(self._entries) >= MAX_KNOWN_CACHED_CONTENTS:
|
||||
self._entries = {k: v for k, v in self._entries.items() if v[1] > now}
|
||||
while len(self._entries) >= MAX_KNOWN_CACHED_CONTENTS:
|
||||
del self._entries[next(iter(self._entries))]
|
||||
self._entries[(scope, cache_key)] = (name, reuse_until)
|
||||
|
||||
def forget_rejected(self, cached_content: str | None, status_code: int) -> None:
|
||||
if cached_content is None or status_code not in _CACHE_REJECTED_STATUS_CODES:
|
||||
return
|
||||
with self._lock:
|
||||
self._entries = {key: entry for key, entry in self._entries.items() if entry[0] != cached_content}
|
||||
|
||||
async def single_flight(
|
||||
self,
|
||||
scope: tuple[str, str],
|
||||
cache_key: str,
|
||||
resolve: Callable[[], Awaitable[str | None]],
|
||||
wait_seconds: float | None,
|
||||
) -> str | None:
|
||||
"""Run ``resolve`` once per (scope, cache_key) at a time within this event loop.
|
||||
|
||||
Concurrent requests for the same prefix wait for the first one's lookup-or-create
|
||||
instead of each listing (and creating a duplicate cache). A waiter resolves on its
|
||||
own when that first attempt fails or has not finished within ``wait_seconds``, and
|
||||
never cancels it.
|
||||
"""
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
own: Final[asyncio.Future[str | None]] = loop.create_future()
|
||||
with self._lock:
|
||||
inflight: Final = self._inflight.setdefault(loop, {})
|
||||
leader: Final = inflight.setdefault((scope, cache_key), own)
|
||||
if leader is not own:
|
||||
with contextlib.suppress(asyncio.TimeoutError):
|
||||
if name := await asyncio.wait_for(asyncio.shield(leader), wait_seconds):
|
||||
return name
|
||||
return await resolve()
|
||||
try:
|
||||
resolved: Final = await resolve()
|
||||
except BaseException:
|
||||
own.set_result(None)
|
||||
raise
|
||||
else:
|
||||
own.set_result(resolved)
|
||||
return resolved
|
||||
finally:
|
||||
with self._lock:
|
||||
inflight.pop((scope, cache_key), None)
|
||||
|
||||
def clear(self) -> None:
|
||||
with self._lock:
|
||||
self._entries.clear()
|
||||
self._inflight.clear()
|
||||
|
||||
|
||||
known_cached_contents: Final = KnownCachedContents()
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
from collections.abc import Awaitable, Callable
|
||||
from copy import deepcopy
|
||||
from functools import partial
|
||||
from typing import Final, TypeVar
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import independent_snapshot
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.llms.vertex_ai import RequestBody
|
||||
from litellm.utils import is_cached_message
|
||||
|
||||
from .known_cached_contents import names_rejected_cache
|
||||
|
||||
Built: Final = TypeVar("Built")
|
||||
Sent: Final = TypeVar("Sent")
|
||||
|
||||
|
||||
def cache_rebuild(
|
||||
build: Callable[..., Built], messages: list[AllMessageValues], optional_params: dict
|
||||
) -> Callable[[], Built] | None:
|
||||
"""Copies the inputs before the first transform, which pops keys from ``optional_params`` and ``messages``."""
|
||||
if not any(map(is_cached_message, messages)):
|
||||
return None
|
||||
return partial(build, messages=deepcopy(messages), optional_params=independent_snapshot(optional_params))
|
||||
|
||||
|
||||
def _status_and_message(error: httpx.HTTPStatusError | BaseLLMException) -> tuple[int, str]:
|
||||
match error:
|
||||
case httpx.HTTPStatusError():
|
||||
return error.response.status_code, error.response.text
|
||||
case BaseLLMException():
|
||||
return error.status_code, error.message
|
||||
|
||||
|
||||
def _cache_rejected(body: RequestBody, error: httpx.HTTPStatusError | BaseLLMException) -> bool:
|
||||
return body.get("cachedContent") is not None and names_rejected_cache(*_status_and_message(error))
|
||||
|
||||
|
||||
async def post_with_fresh_cache_async(
|
||||
body: RequestBody,
|
||||
rebuild: Callable[[], Awaitable[RequestBody]] | None,
|
||||
post: Callable[[RequestBody], Awaitable[Sent]],
|
||||
log_request: Callable[[RequestBody], None],
|
||||
) -> tuple[RequestBody, Sent]:
|
||||
try:
|
||||
return body, await post(body)
|
||||
except (httpx.HTTPStatusError, BaseLLMException) as error:
|
||||
if rebuild is None or not _cache_rejected(body, error):
|
||||
raise
|
||||
fresh_body: Final = await rebuild()
|
||||
log_request(fresh_body)
|
||||
if fresh_body.get("cachedContent") == body.get("cachedContent"):
|
||||
raise
|
||||
return fresh_body, await post(fresh_body)
|
||||
|
||||
|
||||
def post_with_fresh_cache_sync(
|
||||
body: RequestBody,
|
||||
rebuild: Callable[[], RequestBody] | None,
|
||||
post: Callable[[RequestBody], Sent],
|
||||
log_request: Callable[[RequestBody], None],
|
||||
) -> tuple[RequestBody, Sent]:
|
||||
try:
|
||||
return body, post(body)
|
||||
except (httpx.HTTPStatusError, BaseLLMException) as error:
|
||||
if rebuild is None or not _cache_rejected(body, error):
|
||||
raise
|
||||
fresh_body: Final = rebuild()
|
||||
log_request(fresh_body)
|
||||
if fresh_body.get("cachedContent") == body.get("cachedContent"):
|
||||
raise
|
||||
return fresh_body, post(fresh_body)
|
||||
|
|
@ -20,6 +20,7 @@ from litellm.utils import is_prompt_caching_valid_prompt
|
|||
|
||||
from ..common_utils import VertexAIError, get_vertex_base_url
|
||||
from ..vertex_llm_base import VertexBase
|
||||
from .known_cached_contents import KnownCachedContents, known_cached_contents
|
||||
from .transformation import (
|
||||
cached_messages_end_on_supported_turn,
|
||||
separate_cached_messages,
|
||||
|
|
@ -35,6 +36,7 @@ MAX_PAGINATION_PAGES: Final = 100 # Reasonable upper bound for pagination
|
|||
class _CachedContentItem(TypedDict):
|
||||
displayName: ReadOnly[NotRequired[object]]
|
||||
name: ReadOnly[NotRequired[str | None]]
|
||||
expireTime: ReadOnly[NotRequired[str | None]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
|
|
@ -51,6 +53,7 @@ class _CachedContentUsageMetadata(TypedDict):
|
|||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _CreatedCachedContent(TypedDict):
|
||||
name: ReadOnly[NotRequired[str | None]]
|
||||
expireTime: ReadOnly[NotRequired[str | None]]
|
||||
usageMetadata: ReadOnly[NotRequired[_CachedContentUsageMetadata | None]]
|
||||
|
||||
|
||||
|
|
@ -67,6 +70,14 @@ def _record_cache_creation_tokens(logging_obj: Logging, created: _CreatedCachedC
|
|||
logging_obj.vertex_ai_cache_creation_input_tokens = max(tokens, 0) # rebind-ok: set provider count
|
||||
|
||||
|
||||
def _wait_seconds(timeout: float | httpx.Timeout | None) -> float | None:
|
||||
match timeout:
|
||||
case httpx.Timeout():
|
||||
return timeout.read
|
||||
case _:
|
||||
return timeout
|
||||
|
||||
|
||||
class ContextCachingEndpoints(VertexBase):
|
||||
"""
|
||||
Covers context caching endpoints for Vertex AI + Google AI Studio
|
||||
|
|
@ -74,8 +85,9 @@ class ContextCachingEndpoints(VertexBase):
|
|||
v0: covers Google AI Studio
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
def __init__(self, known_caches: KnownCachedContents = known_cached_contents) -> None:
|
||||
super().__init__()
|
||||
self._known_caches: Final = known_caches
|
||||
|
||||
def _get_token_and_url_context_caching(
|
||||
self,
|
||||
|
|
@ -129,7 +141,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
self,
|
||||
cache_key: str,
|
||||
client: HTTPHandler,
|
||||
headers: dict,
|
||||
headers: dict[str, str], # mutable-ok: passed through to the HTTP client, which takes a dict
|
||||
api_key: str,
|
||||
api_base: str | None,
|
||||
logging_obj: Logging,
|
||||
|
|
@ -159,6 +171,12 @@ class ContextCachingEndpoints(VertexBase):
|
|||
vertex_auth_header=vertex_auth_header,
|
||||
model=model,
|
||||
)
|
||||
scope: Final = KnownCachedContents.scope(
|
||||
base_url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location
|
||||
)
|
||||
known_name: Final = self._known_caches.get(scope, cache_key)
|
||||
if known_name is not None:
|
||||
return known_name
|
||||
|
||||
page_token: str | None = None
|
||||
|
||||
|
|
@ -201,8 +219,14 @@ class ContextCachingEndpoints(VertexBase):
|
|||
# Check current page for matching cache_key
|
||||
for cached_item in raw_response["cachedContents"]:
|
||||
display_name = cached_item.get("displayName")
|
||||
if display_name is not None and display_name == cache_key:
|
||||
return cached_item.get("name")
|
||||
if display_name is None or display_name != cache_key:
|
||||
continue
|
||||
expire_time = cached_item.get("expireTime")
|
||||
if not self._known_caches.has_not_expired(expire_time):
|
||||
continue
|
||||
name = cached_item.get("name")
|
||||
self._known_caches.remember(scope, cache_key, name, expire_time)
|
||||
return name
|
||||
|
||||
# Check if there are more pages
|
||||
page_token = raw_response.get("nextPageToken")
|
||||
|
|
@ -216,7 +240,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
self,
|
||||
cache_key: str,
|
||||
client: AsyncHTTPHandler,
|
||||
headers: dict,
|
||||
headers: dict[str, str], # mutable-ok: passed through to the HTTP client, which takes a dict
|
||||
api_key: str,
|
||||
api_base: str | None,
|
||||
logging_obj: Logging,
|
||||
|
|
@ -246,6 +270,12 @@ class ContextCachingEndpoints(VertexBase):
|
|||
vertex_auth_header=vertex_auth_header,
|
||||
model=model,
|
||||
)
|
||||
scope: Final = KnownCachedContents.scope(
|
||||
base_url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location
|
||||
)
|
||||
known_name: Final = self._known_caches.get(scope, cache_key)
|
||||
if known_name is not None:
|
||||
return known_name
|
||||
|
||||
page_token: str | None = None
|
||||
|
||||
|
|
@ -288,8 +318,14 @@ class ContextCachingEndpoints(VertexBase):
|
|||
# Check current page for matching cache_key
|
||||
for cached_item in raw_response["cachedContents"]:
|
||||
display_name = cached_item.get("displayName")
|
||||
if display_name is not None and display_name == cache_key:
|
||||
return cached_item.get("name")
|
||||
if display_name is None or display_name != cache_key:
|
||||
continue
|
||||
expire_time = cached_item.get("expireTime")
|
||||
if not self._known_caches.has_not_expired(expire_time):
|
||||
continue
|
||||
name = cached_item.get("name")
|
||||
self._known_caches.remember(scope, cache_key, name, expire_time)
|
||||
return name
|
||||
|
||||
# Check if there are more pages
|
||||
page_token = raw_response.get("nextPageToken")
|
||||
|
|
@ -451,6 +487,14 @@ class ContextCachingEndpoints(VertexBase):
|
|||
|
||||
raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root
|
||||
_record_cache_creation_tokens(logging_obj, raw_response_cached)
|
||||
self._known_caches.remember(
|
||||
KnownCachedContents.scope(
|
||||
url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location
|
||||
),
|
||||
generated_cache_key,
|
||||
raw_response_cached.get("name"),
|
||||
raw_response_cached.get("expireTime"),
|
||||
)
|
||||
return (
|
||||
non_cached_messages,
|
||||
optional_params,
|
||||
|
|
@ -548,68 +592,77 @@ class ContextCachingEndpoints(VertexBase):
|
|||
generated_cache_key: Final = local_cache_obj.get_cache_key(
|
||||
messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model
|
||||
)
|
||||
google_cache_name: Final = await self.async_check_cache(
|
||||
cache_key=generated_cache_key,
|
||||
client=client,
|
||||
headers=headers,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
vertex_auth_header=vertex_auth_header,
|
||||
model=model,
|
||||
scope: Final = KnownCachedContents.scope(
|
||||
url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location
|
||||
)
|
||||
|
||||
if google_cache_name:
|
||||
return non_cached_messages, optional_params, google_cache_name
|
||||
|
||||
## TRANSFORM REQUEST
|
||||
cached_content_request_body: Final = transform_openai_messages_to_gemini_context_caching(
|
||||
model=model,
|
||||
messages=cached_messages,
|
||||
cache_key=generated_cache_key,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
cached_content_request_body["tools"] = tools
|
||||
if tool_choice is not None:
|
||||
cached_content_request_body["toolConfig"] = tool_choice
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=messages,
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": cached_content_request_body,
|
||||
"api_base": url,
|
||||
"headers": headers,
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
response: Final = await client.post(
|
||||
url=url,
|
||||
async def lookup_or_create() -> str | None:
|
||||
google_cache_name: Final = await self.async_check_cache(
|
||||
cache_key=generated_cache_key,
|
||||
client=client,
|
||||
headers=headers,
|
||||
json=cached_content_request_body,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
vertex_auth_header=vertex_auth_header,
|
||||
model=model,
|
||||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
error_code: Final = err.response.status_code
|
||||
raise VertexAIError(status_code=error_code, message=err.response.text)
|
||||
except httpx.TimeoutException:
|
||||
raise VertexAIError(status_code=408, message="Timeout error occurred.")
|
||||
|
||||
raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root
|
||||
_record_cache_creation_tokens(logging_obj, raw_response_cached)
|
||||
return (
|
||||
non_cached_messages,
|
||||
optional_params,
|
||||
raw_response_cached.get("name"),
|
||||
if google_cache_name:
|
||||
return google_cache_name
|
||||
|
||||
## TRANSFORM REQUEST
|
||||
cached_content_request_body: Final = transform_openai_messages_to_gemini_context_caching(
|
||||
model=model,
|
||||
messages=cached_messages,
|
||||
cache_key=generated_cache_key,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
cached_content_request_body["tools"] = tools
|
||||
if tool_choice is not None:
|
||||
cached_content_request_body["toolConfig"] = tool_choice
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=messages,
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": cached_content_request_body,
|
||||
"api_base": url,
|
||||
"headers": headers,
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
response: Final = await client.post(
|
||||
url=url,
|
||||
headers=headers,
|
||||
json=cached_content_request_body,
|
||||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
error_code: Final = err.response.status_code
|
||||
raise VertexAIError(status_code=error_code, message=err.response.text)
|
||||
except httpx.TimeoutException:
|
||||
raise VertexAIError(status_code=408, message="Timeout error occurred.")
|
||||
|
||||
raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root
|
||||
_record_cache_creation_tokens(logging_obj, raw_response_cached)
|
||||
self._known_caches.remember(
|
||||
scope, generated_cache_key, raw_response_cached.get("name"), raw_response_cached.get("expireTime")
|
||||
)
|
||||
return raw_response_cached.get("name")
|
||||
|
||||
cache_name: Final = await self._known_caches.single_flight(
|
||||
scope, generated_cache_key, lookup_or_create, _wait_seconds(timeout)
|
||||
)
|
||||
return non_cached_messages, optional_params, cache_name
|
||||
|
||||
def get_cache(self):
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -96,6 +96,12 @@ from ..common_utils import (
|
|||
build_vertex_schema,
|
||||
supports_response_json_schema,
|
||||
)
|
||||
from ..context_caching.known_cached_contents import known_cached_contents
|
||||
from ..context_caching.rejected_cache_retry import (
|
||||
cache_rebuild,
|
||||
post_with_fresh_cache_async,
|
||||
post_with_fresh_cache_sync,
|
||||
)
|
||||
from ..vertex_llm_base import VertexBase
|
||||
from .grounding_requests import calculate_grounding_requests
|
||||
from .long_decode import (
|
||||
|
|
@ -2786,7 +2792,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
litellm_params: dict,
|
||||
api_key: str | dict | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict:
|
||||
) -> dict[str, str]:
|
||||
default_headers: Final = {
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
|
@ -2800,6 +2806,20 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
return default_headers
|
||||
|
||||
|
||||
def _log_generate_request(
|
||||
logging_obj: LoggingClass,
|
||||
messages: Sequence[AllMessageValues],
|
||||
api_base: str,
|
||||
headers: Mapping[str, str],
|
||||
body: RequestBody,
|
||||
) -> None:
|
||||
logging_obj.pre_call(
|
||||
input=messages,
|
||||
api_key="",
|
||||
additional_args={"complete_input_dict": body, "api_base": api_base, "headers": headers},
|
||||
)
|
||||
|
||||
|
||||
async def make_call(
|
||||
client: AsyncHTTPHandler | None, # module-level client
|
||||
gemini_client: AsyncHTTPHandler | None, # if passed by user
|
||||
|
|
@ -2832,6 +2852,7 @@ async def make_call(
|
|||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as e:
|
||||
exception_string: Final = str(await e.response.aread())
|
||||
known_cached_contents.forget_rejected(request_body.get("cachedContent"), e.response.status_code)
|
||||
raise VertexAIError(
|
||||
status_code=e.response.status_code,
|
||||
message=VertexGeminiConfig().translate_exception_str(exception_string),
|
||||
|
|
@ -2845,14 +2866,18 @@ async def make_call(
|
|||
)
|
||||
|
||||
async def open_hop(body: RequestBody) -> httpx.Response:
|
||||
return await client.post(
|
||||
url_for(api_base, body),
|
||||
headers=dict(await hop_headers()),
|
||||
data=json.dumps(body),
|
||||
stream=True,
|
||||
logging_obj=logging_obj,
|
||||
timeout=timeout,
|
||||
)
|
||||
try:
|
||||
return await client.post(
|
||||
url_for(api_base, body),
|
||||
headers=dict(await hop_headers()),
|
||||
data=json.dumps(body),
|
||||
stream=True,
|
||||
logging_obj=logging_obj,
|
||||
timeout=timeout,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
known_cached_contents.forget_rejected(body.get("cachedContent"), e.response.status_code)
|
||||
raise
|
||||
|
||||
completion_stream: Final = ModelResponseIterator(
|
||||
streaming_response=response.aiter_lines(),
|
||||
|
|
@ -2891,14 +2916,18 @@ def make_sync_call(
|
|||
if client is None:
|
||||
client = HTTPHandler() # Create a new client if none provided
|
||||
|
||||
response: Final = client.post(
|
||||
url_for(api_base, request_body),
|
||||
headers=headers,
|
||||
data=data,
|
||||
stream=True,
|
||||
logging_obj=logging_obj,
|
||||
timeout=timeout,
|
||||
)
|
||||
try:
|
||||
response: Final = client.post(
|
||||
url_for(api_base, request_body),
|
||||
headers=headers,
|
||||
data=data,
|
||||
stream=True,
|
||||
logging_obj=logging_obj,
|
||||
timeout=timeout,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
known_cached_contents.forget_rejected(request_body.get("cachedContent"), e.response.status_code)
|
||||
raise
|
||||
|
||||
if response.status_code != 200 and response.status_code != 201:
|
||||
raise VertexAIError(
|
||||
|
|
@ -2908,14 +2937,18 @@ def make_sync_call(
|
|||
)
|
||||
|
||||
def open_hop(body: RequestBody) -> httpx.Response:
|
||||
return client.post(
|
||||
url_for(api_base, body),
|
||||
headers=dict(hop_headers()),
|
||||
data=json.dumps(body),
|
||||
stream=True,
|
||||
logging_obj=logging_obj,
|
||||
timeout=timeout,
|
||||
)
|
||||
try:
|
||||
return client.post(
|
||||
url_for(api_base, body),
|
||||
headers=dict(hop_headers()),
|
||||
data=json.dumps(body),
|
||||
stream=True,
|
||||
logging_obj=logging_obj,
|
||||
timeout=timeout,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
known_cached_contents.forget_rejected(body.get("cachedContent"), e.response.status_code)
|
||||
raise
|
||||
|
||||
completion_stream: Final = ModelResponseIterator(
|
||||
streaming_response=response.iter_lines(),
|
||||
|
|
@ -2983,14 +3016,14 @@ class VertexLLM(VertexBase):
|
|||
custom_llm_provider: Literal[
|
||||
"vertex_ai", "vertex_ai_beta", "gemini"
|
||||
], # if it's vertex_ai or gemini (google ai studio)
|
||||
messages: list,
|
||||
messages: list[AllMessageValues],
|
||||
model_response: ModelResponse,
|
||||
print_verbose: Callable,
|
||||
data: dict,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
encoding: object,
|
||||
logging_obj,
|
||||
stream,
|
||||
logging_obj: LoggingClass,
|
||||
stream: bool | None,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
logger_fn=None,
|
||||
|
|
@ -3033,6 +3066,17 @@ class VertexLLM(VertexBase):
|
|||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
rebuild: Final = cache_rebuild(
|
||||
partial(
|
||||
async_transform_request_body,
|
||||
**data,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
vertex_auth_header=auth_header,
|
||||
),
|
||||
messages,
|
||||
optional_params,
|
||||
)
|
||||
request_body: Final = await async_transform_request_body(
|
||||
**data,
|
||||
vertex_project=vertex_project,
|
||||
|
|
@ -3041,39 +3085,37 @@ class VertexLLM(VertexBase):
|
|||
)
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=messages,
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": request_body,
|
||||
"api_base": api_base,
|
||||
"headers": headers,
|
||||
},
|
||||
log_request: Final = partial(_log_generate_request, logging_obj, messages, api_base, headers)
|
||||
log_request(request_body)
|
||||
|
||||
open_stream: Final = partial(
|
||||
make_call,
|
||||
gemini_client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None),
|
||||
api_base=api_base,
|
||||
headers=headers,
|
||||
model=model,
|
||||
messages=messages,
|
||||
logging_obj=logging_obj,
|
||||
hop_headers=partial(
|
||||
self._fresh_headers_async,
|
||||
headers,
|
||||
extra_headers,
|
||||
vertex_credentials,
|
||||
vertex_project,
|
||||
custom_llm_provider,
|
||||
),
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
request_body_str: Final = json.dumps(request_body)
|
||||
async def open_with_fresh_cache(client: AsyncHTTPHandler | None) -> ModelResponseIterator:
|
||||
async def post(body: RequestBody) -> ModelResponseIterator:
|
||||
return await open_stream(client=client, data=json.dumps(body), request_body=body)
|
||||
|
||||
return (await post_with_fresh_cache_async(request_body, rebuild, post, log_request))[1]
|
||||
|
||||
streaming_response: Final = CustomStreamWrapper(
|
||||
completion_stream=None,
|
||||
make_call=partial(
|
||||
make_call,
|
||||
gemini_client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None),
|
||||
api_base=api_base,
|
||||
headers=headers,
|
||||
data=request_body_str,
|
||||
model=model,
|
||||
messages=messages,
|
||||
logging_obj=logging_obj,
|
||||
request_body=request_body,
|
||||
hop_headers=partial(
|
||||
self._fresh_headers_async,
|
||||
headers,
|
||||
extra_headers,
|
||||
vertex_credentials,
|
||||
vertex_project,
|
||||
custom_llm_provider,
|
||||
),
|
||||
timeout=timeout,
|
||||
),
|
||||
make_call=open_with_fresh_cache,
|
||||
model=model,
|
||||
custom_llm_provider="vertex_ai_beta",
|
||||
logging_obj=logging_obj,
|
||||
|
|
@ -3083,7 +3125,7 @@ class VertexLLM(VertexBase):
|
|||
async def async_completion(
|
||||
self,
|
||||
model: str,
|
||||
messages: list,
|
||||
messages: list[AllMessageValues],
|
||||
model_response: ModelResponse,
|
||||
print_verbose: Callable,
|
||||
data: dict,
|
||||
|
|
@ -3093,7 +3135,7 @@ class VertexLLM(VertexBase):
|
|||
timeout: float | httpx.Timeout | None,
|
||||
encoding: object,
|
||||
logging_obj: LoggingClass,
|
||||
stream,
|
||||
stream: bool | None,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
logger_fn=None,
|
||||
|
|
@ -3136,6 +3178,17 @@ class VertexLLM(VertexBase):
|
|||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
rebuild: Final = cache_rebuild(
|
||||
partial(
|
||||
async_transform_request_body,
|
||||
**data,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
vertex_auth_header=auth_header,
|
||||
),
|
||||
messages,
|
||||
optional_params,
|
||||
)
|
||||
request_body: Final = await async_transform_request_body(
|
||||
**data,
|
||||
vertex_project=vertex_project,
|
||||
|
|
@ -3151,15 +3204,8 @@ class VertexLLM(VertexBase):
|
|||
else:
|
||||
client = client
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=messages,
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": request_body,
|
||||
"api_base": api_base,
|
||||
"headers": headers,
|
||||
},
|
||||
)
|
||||
log_request: Final = partial(_log_generate_request, logging_obj, messages, api_base, headers)
|
||||
log_request(request_body)
|
||||
|
||||
fresh_headers: Final = partial(
|
||||
self._fresh_headers_async, headers, extra_headers, vertex_credentials, vertex_project, custom_llm_provider
|
||||
|
|
@ -3175,6 +3221,7 @@ class VertexLLM(VertexBase):
|
|||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
known_cached_contents.forget_rejected(body.get("cachedContent"), err.response.status_code)
|
||||
raise VertexAIError(
|
||||
status_code=err.response.status_code,
|
||||
message=err.response.text,
|
||||
|
|
@ -3188,14 +3235,14 @@ class VertexLLM(VertexBase):
|
|||
)
|
||||
return response
|
||||
|
||||
first_response: Final = await post(request_body)
|
||||
sent_body, first_response = await post_with_fresh_cache_async(request_body, rebuild, post, log_request)
|
||||
return _transform_long_decode_hops(
|
||||
model=model,
|
||||
hops=await collect_hops_async(parse_body(first_response), request_body, post),
|
||||
hops=await collect_hops_async(parse_body(first_response), sent_body, post),
|
||||
first_response=first_response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
request_body=request_body,
|
||||
request_body=sent_body,
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
|
|
@ -3320,6 +3367,17 @@ class VertexLLM(VertexBase):
|
|||
)
|
||||
|
||||
## TRANSFORMATION ##
|
||||
rebuild: Final = cache_rebuild(
|
||||
partial(
|
||||
sync_transform_request_body,
|
||||
**transform_request_params,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
vertex_auth_header=auth_header,
|
||||
),
|
||||
messages,
|
||||
optional_params,
|
||||
)
|
||||
data: Final = sync_transform_request_body(
|
||||
**transform_request_params,
|
||||
vertex_project=vertex_project,
|
||||
|
|
@ -3328,15 +3386,8 @@ class VertexLLM(VertexBase):
|
|||
)
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=messages,
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": data,
|
||||
"api_base": url,
|
||||
"headers": headers,
|
||||
},
|
||||
)
|
||||
log_request: Final = partial(_log_generate_request, logging_obj, messages, url, headers)
|
||||
log_request(data)
|
||||
|
||||
fresh_headers: Final = partial(
|
||||
self._fresh_headers, headers, extra_headers, vertex_credentials, vertex_project, custom_llm_provider
|
||||
|
|
@ -3344,22 +3395,27 @@ class VertexLLM(VertexBase):
|
|||
|
||||
## SYNC STREAMING CALL ##
|
||||
if stream is True:
|
||||
request_data_str: Final = json.dumps(data)
|
||||
open_stream: Final = partial(
|
||||
make_sync_call,
|
||||
gemini_client=(client if client is not None and isinstance(client, HTTPHandler) else None),
|
||||
api_base=url,
|
||||
model=model,
|
||||
messages=messages,
|
||||
logging_obj=logging_obj,
|
||||
headers=headers,
|
||||
hop_headers=fresh_headers,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
def open_with_fresh_cache(client: HTTPHandler | None) -> ModelResponseIterator:
|
||||
def post(body: RequestBody) -> ModelResponseIterator:
|
||||
return open_stream(client=client, data=json.dumps(body), request_body=body)
|
||||
|
||||
return post_with_fresh_cache_sync(data, rebuild, post, log_request)[1]
|
||||
|
||||
streaming_response: Final = CustomStreamWrapper(
|
||||
completion_stream=None,
|
||||
make_call=partial(
|
||||
make_sync_call,
|
||||
gemini_client=(client if client is not None and isinstance(client, HTTPHandler) else None),
|
||||
api_base=url,
|
||||
data=request_data_str,
|
||||
model=model,
|
||||
messages=messages,
|
||||
logging_obj=logging_obj,
|
||||
headers=headers,
|
||||
request_body=data,
|
||||
hop_headers=fresh_headers,
|
||||
timeout=timeout,
|
||||
),
|
||||
make_call=open_with_fresh_cache,
|
||||
model=model,
|
||||
custom_llm_provider="vertex_ai_beta",
|
||||
logging_obj=logging_obj,
|
||||
|
|
@ -3385,6 +3441,7 @@ class VertexLLM(VertexBase):
|
|||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
known_cached_contents.forget_rejected(body.get("cachedContent"), err.response.status_code)
|
||||
raise VertexAIError(
|
||||
status_code=err.response.status_code,
|
||||
message=err.response.text,
|
||||
|
|
@ -3398,14 +3455,14 @@ class VertexLLM(VertexBase):
|
|||
)
|
||||
return response
|
||||
|
||||
first_response: Final = post(data)
|
||||
sent_body, first_response = post_with_fresh_cache_sync(data, rebuild, post, log_request)
|
||||
return _transform_long_decode_hops(
|
||||
model=model,
|
||||
hops=collect_hops_sync(parse_body(first_response), data, post),
|
||||
hops=collect_hops_sync(parse_body(first_response), sent_body, post),
|
||||
first_response=first_response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
request_body=data,
|
||||
request_body=sent_body,
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
|
|
|
|||
10
tests/unit/llms/vertex_ai/context_caching/conftest.py
Normal file
10
tests/unit/llms/vertex_ai/context_caching/conftest.py
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
import pytest
|
||||
|
||||
from litellm.llms.vertex_ai.context_caching.known_cached_contents import known_cached_contents
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _forget_known_cached_contents():
|
||||
known_cached_contents.clear()
|
||||
yield
|
||||
known_cached_contents.clear()
|
||||
|
|
@ -0,0 +1,663 @@
|
|||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
from litellm.llms.vertex_ai.context_caching import vertex_ai_context_caching
|
||||
from litellm.llms.vertex_ai.context_caching.known_cached_contents import KnownCachedContents, parse_expire_time
|
||||
from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ContextCachingEndpoints
|
||||
from litellm.utils import get_prompt_cache_min_tokens
|
||||
|
||||
_REGISTRY_MODULE: Final = "litellm.llms.vertex_ai.context_caching.known_cached_contents"
|
||||
_SCOPE: Final = ("https://vertex.test/v1/projects/p/locations/us/cachedContents", "")
|
||||
_NOW: Final = datetime(2026, 10, 14, 18, 0, tzinfo=timezone.utc).timestamp()
|
||||
_FAR_FUTURE: Final = "2100-01-01T00:00:00.123456789Z"
|
||||
_GATEWAY: Final = "https://gateway.test/vertex"
|
||||
_PREFIX_REPEATS: Final = max(
|
||||
get_prompt_cache_min_tokens(model=f"{provider}/gemini-3.1-pro") for provider in ("vertex_ai", "gemini")
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clock() -> MagicMock:
|
||||
return MagicMock(return_value=_NOW)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def endpoints(clock: MagicMock) -> ContextCachingEndpoints:
|
||||
return ContextCachingEndpoints(known_caches=KnownCachedContents(clock=clock))
|
||||
|
||||
|
||||
def _expire_time(seconds_from_now: float) -> str:
|
||||
expires: Final = datetime.fromtimestamp(_NOW + seconds_from_now, timezone.utc)
|
||||
return expires.strftime("%Y-%m-%dT%H:%M:%S.%f") + "123Z"
|
||||
|
||||
|
||||
def _messages(label: str = "shared") -> list:
|
||||
text: Final = f"{label} cacheable prefix " * _PREFIX_REPEATS
|
||||
return [
|
||||
{"role": "system", "content": [{"type": "text", "text": text, "cache_control": {"type": "ephemeral"}}]},
|
||||
{"role": "user", "content": "question"},
|
||||
]
|
||||
|
||||
|
||||
def _response(body: dict) -> MagicMock:
|
||||
return MagicMock(json=lambda: body)
|
||||
|
||||
|
||||
def _created(name: str, expire_time: str | None) -> dict:
|
||||
body: Final = {"name": name, "usageMetadata": {"totalTokenCount": 2048}}
|
||||
return body if expire_time is None else {**body, "expireTime": expire_time}
|
||||
|
||||
|
||||
def _sync_call(
|
||||
endpoints: ContextCachingEndpoints,
|
||||
client: HTTPHandler,
|
||||
api_key: str = "k",
|
||||
custom_llm_provider: str = "vertex_ai",
|
||||
messages=None,
|
||||
extra_headers: dict | None = None,
|
||||
api_base: str | None = None,
|
||||
vertex_project: str = "p",
|
||||
):
|
||||
return endpoints.check_and_create_cache(
|
||||
extra_headers=extra_headers,
|
||||
messages=messages or _messages(),
|
||||
optional_params={},
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
model="gemini-3.1-pro",
|
||||
client=client,
|
||||
timeout=30.0,
|
||||
logging_obj=MagicMock(model_call_details={}),
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location="us",
|
||||
vertex_auth_header="token",
|
||||
)
|
||||
|
||||
|
||||
async def _async_call(
|
||||
endpoints: ContextCachingEndpoints,
|
||||
client: AsyncHTTPHandler,
|
||||
timeout: float | httpx.Timeout = 30.0,
|
||||
api_base: str | None = None,
|
||||
vertex_project: str = "p",
|
||||
):
|
||||
return await endpoints.async_check_and_create_cache(
|
||||
messages=_messages(),
|
||||
optional_params={},
|
||||
api_key="k",
|
||||
api_base=api_base,
|
||||
model="gemini-3.1-pro",
|
||||
client=client,
|
||||
timeout=timeout,
|
||||
logging_obj=MagicMock(model_call_details={}),
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_project=vertex_project,
|
||||
vertex_location="us",
|
||||
vertex_auth_header="token",
|
||||
)
|
||||
|
||||
|
||||
def test_created_cache_is_reused_without_listing(endpoints: ContextCachingEndpoints):
|
||||
client: Final = MagicMock(spec=HTTPHandler)
|
||||
client.get.return_value = _response({})
|
||||
client.post.return_value = _response(_created("cachedContents/1", _expire_time(60)))
|
||||
|
||||
assert _sync_call(endpoints, client)[2] == "cachedContents/1"
|
||||
assert _sync_call(endpoints, client)[2] == "cachedContents/1"
|
||||
|
||||
assert client.get.call_count == 1
|
||||
assert client.post.call_count == 1
|
||||
|
||||
|
||||
def test_listed_cache_is_reused_without_listing_again(endpoints: ContextCachingEndpoints):
|
||||
client: Final = MagicMock(spec=HTTPHandler)
|
||||
key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key(
|
||||
messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro"
|
||||
)
|
||||
client.get.return_value = _response(
|
||||
{"cachedContents": [{"name": "cachedContents/listed", "displayName": key, "expireTime": _expire_time(60)}]}
|
||||
)
|
||||
|
||||
assert _sync_call(endpoints, client)[2] == "cachedContents/listed"
|
||||
assert _sync_call(endpoints, client)[2] == "cachedContents/listed"
|
||||
|
||||
assert client.get.call_count == 1
|
||||
client.post.assert_not_called()
|
||||
|
||||
|
||||
def test_listed_cache_past_its_expire_time_is_not_used(endpoints: ContextCachingEndpoints):
|
||||
client: Final = MagicMock(spec=HTTPHandler)
|
||||
key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key(
|
||||
messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro"
|
||||
)
|
||||
client.get.return_value = _response(
|
||||
{"cachedContents": [{"name": "cachedContents/stale", "displayName": key, "expireTime": _expire_time(-1)}]}
|
||||
)
|
||||
client.post.return_value = _response(_created("cachedContents/fresh", _expire_time(60)))
|
||||
|
||||
assert _sync_call(endpoints, client)[2] == "cachedContents/fresh"
|
||||
|
||||
|
||||
def test_short_lived_cache_is_reused_from_the_listing(endpoints: ContextCachingEndpoints, clock: MagicMock):
|
||||
client: Final = MagicMock(spec=HTTPHandler)
|
||||
key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key(
|
||||
messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro"
|
||||
)
|
||||
client.get.side_effect = [
|
||||
_response({}),
|
||||
_response(
|
||||
{"cachedContents": [{"name": "cachedContents/1", "displayName": key, "expireTime": _expire_time(8)}]}
|
||||
),
|
||||
]
|
||||
client.post.return_value = _response(_created("cachedContents/1", _expire_time(8)))
|
||||
|
||||
assert _sync_call(endpoints, client)[2] == "cachedContents/1"
|
||||
clock.return_value = _NOW + 1
|
||||
assert _sync_call(endpoints, client)[2] == "cachedContents/1"
|
||||
|
||||
assert client.post.call_count == 1
|
||||
|
||||
|
||||
def test_cache_without_expire_time_is_not_remembered(endpoints: ContextCachingEndpoints):
|
||||
client: Final = MagicMock(spec=HTTPHandler)
|
||||
client.get.return_value = _response({})
|
||||
client.post.return_value = _response(_created("cachedContents/1", None))
|
||||
|
||||
_sync_call(endpoints, client)
|
||||
_sync_call(endpoints, client)
|
||||
|
||||
assert client.get.call_count == 2
|
||||
|
||||
|
||||
def test_remembered_cache_is_listed_again_once_its_reuse_window_closes(
|
||||
endpoints: ContextCachingEndpoints, clock: MagicMock
|
||||
):
|
||||
client: Final = MagicMock(spec=HTTPHandler)
|
||||
key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key(
|
||||
messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro"
|
||||
)
|
||||
client.get.side_effect = [
|
||||
_response({}),
|
||||
_response(
|
||||
{"cachedContents": [{"name": "cachedContents/listed", "displayName": key, "expireTime": _expire_time(600)}]}
|
||||
),
|
||||
]
|
||||
client.post.return_value = _response(_created("cachedContents/1", _expire_time(60)))
|
||||
assert _sync_call(endpoints, client)[2] == "cachedContents/1"
|
||||
|
||||
clock.return_value = _NOW + 55
|
||||
assert _sync_call(endpoints, client)[2] == "cachedContents/listed"
|
||||
|
||||
|
||||
def test_gemini_caches_are_not_shared_across_api_keys(endpoints: ContextCachingEndpoints):
|
||||
client: Final = MagicMock(spec=HTTPHandler)
|
||||
client.get.return_value = _response({})
|
||||
client.post.return_value = _response(_created("cachedContents/1", _expire_time(60)))
|
||||
|
||||
_sync_call(endpoints, client, api_key="key-a", custom_llm_provider="gemini")
|
||||
_sync_call(endpoints, client, api_key="key-b", custom_llm_provider="gemini")
|
||||
|
||||
assert client.post.call_count == 2
|
||||
|
||||
|
||||
def test_gemini_cache_created_with_an_overridden_api_key_is_not_reused_with_the_configured_key(
|
||||
endpoints: ContextCachingEndpoints,
|
||||
):
|
||||
client: Final = MagicMock(spec=HTTPHandler)
|
||||
client.get.return_value = _response({})
|
||||
client.post.side_effect = [
|
||||
_response(_created("cachedContents/caller-project", _expire_time(60))),
|
||||
_response(_created("cachedContents/deployment-project", _expire_time(60))),
|
||||
]
|
||||
|
||||
caller: Final = _sync_call(
|
||||
endpoints,
|
||||
client,
|
||||
api_key="deployment-key",
|
||||
custom_llm_provider="gemini",
|
||||
extra_headers={"X-Goog-Api-Key": "caller-key"},
|
||||
)
|
||||
configured: Final = _sync_call(endpoints, client, api_key="deployment-key", custom_llm_provider="gemini")
|
||||
|
||||
assert (caller[2], configured[2]) == ("cachedContents/caller-project", "cachedContents/deployment-project")
|
||||
|
||||
|
||||
def test_vertex_caches_behind_one_custom_api_base_are_not_shared_across_projects(endpoints: ContextCachingEndpoints):
|
||||
client: Final = MagicMock(spec=HTTPHandler)
|
||||
client.get.return_value = _response({})
|
||||
client.post.side_effect = [
|
||||
_response(_created("cachedContents/project-a", _expire_time(60))),
|
||||
_response(_created("cachedContents/project-b", _expire_time(60))),
|
||||
]
|
||||
|
||||
sent: Final = [
|
||||
_sync_call(endpoints, client, api_base=_GATEWAY, vertex_project=project)[2]
|
||||
for project in ("project-a", "project-b", "project-a")
|
||||
]
|
||||
|
||||
assert sent == ["cachedContents/project-a", "cachedContents/project-b", "cachedContents/project-a"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_vertex_caches_behind_one_custom_api_base_are_not_shared_across_projects(
|
||||
endpoints: ContextCachingEndpoints,
|
||||
):
|
||||
client: Final = MagicMock(spec=AsyncHTTPHandler)
|
||||
client.get = AsyncMock(return_value=_response({}))
|
||||
client.post = AsyncMock(
|
||||
side_effect=[
|
||||
_response(_created("cachedContents/project-a", _expire_time(60))),
|
||||
_response(_created("cachedContents/project-b", _expire_time(60))),
|
||||
]
|
||||
)
|
||||
|
||||
sent: Final = [
|
||||
(await _async_call(endpoints, client, api_base=_GATEWAY, vertex_project=project))[2]
|
||||
for project in ("project-a", "project-b", "project-a")
|
||||
]
|
||||
|
||||
assert sent == ["cachedContents/project-a", "cachedContents/project-b", "cachedContents/project-a"]
|
||||
|
||||
|
||||
def test_different_prefixes_do_not_share_a_cache(endpoints: ContextCachingEndpoints):
|
||||
client: Final = MagicMock(spec=HTTPHandler)
|
||||
client.get.return_value = _response({})
|
||||
client.post.side_effect = [
|
||||
_response(_created("cachedContents/a", _expire_time(60))),
|
||||
_response(_created("cachedContents/b", _expire_time(60))),
|
||||
]
|
||||
|
||||
assert _sync_call(endpoints, client, messages=_messages("first"))[2] == "cachedContents/a"
|
||||
assert _sync_call(endpoints, client, messages=_messages("second"))[2] == "cachedContents/b"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_requests_for_one_prefix_list_and_create_once(endpoints: ContextCachingEndpoints):
|
||||
client: Final = MagicMock(spec=AsyncHTTPHandler)
|
||||
|
||||
async def slow_list(**_):
|
||||
await asyncio.sleep(0.01)
|
||||
return _response({})
|
||||
|
||||
client.get = AsyncMock(side_effect=slow_list)
|
||||
client.post = AsyncMock(return_value=_response(_created("cachedContents/1", _expire_time(60))))
|
||||
|
||||
results: Final = await asyncio.gather(*(_async_call(endpoints, client) for _ in range(5)))
|
||||
|
||||
assert {result[2] for result in results} == {"cachedContents/1"}
|
||||
assert client.get.await_count == 1
|
||||
assert client.post.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_created_cache_is_reused_by_a_later_request(endpoints: ContextCachingEndpoints):
|
||||
client: Final = MagicMock(spec=AsyncHTTPHandler)
|
||||
client.get = AsyncMock(return_value=_response({}))
|
||||
client.post = AsyncMock(return_value=_response(_created("cachedContents/1", _expire_time(60))))
|
||||
|
||||
first: Final = await _async_call(endpoints, client)
|
||||
client.get = AsyncMock(side_effect=AssertionError("a known cache must not be listed again"))
|
||||
second: Final = await _async_call(endpoints, client)
|
||||
|
||||
assert (first[2], second[2]) == ("cachedContents/1", "cachedContents/1")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_listed_cache_past_its_expire_time_is_not_used(endpoints: ContextCachingEndpoints):
|
||||
client: Final = MagicMock(spec=AsyncHTTPHandler)
|
||||
key: Final = vertex_ai_context_caching.local_cache_obj.get_cache_key(
|
||||
messages=_messages()[:1], tools=None, tool_choice=None, model="gemini-3.1-pro"
|
||||
)
|
||||
client.get = AsyncMock(
|
||||
return_value=_response(
|
||||
{"cachedContents": [{"name": "cachedContents/stale", "displayName": key, "expireTime": _expire_time(-1)}]}
|
||||
)
|
||||
)
|
||||
client.post = AsyncMock(return_value=_response(_created("cachedContents/fresh", _expire_time(60))))
|
||||
|
||||
assert (await _async_call(endpoints, client))[2] == "cachedContents/fresh"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_create_http_error_surfaces_its_status_code(endpoints: ContextCachingEndpoints):
|
||||
client: Final = MagicMock(spec=AsyncHTTPHandler)
|
||||
client.get = AsyncMock(return_value=_response({}))
|
||||
request: Final = httpx.Request("POST", "https://vertex.test/v1/projects/p/locations/us/cachedContents")
|
||||
failed: Final = httpx.Response(429, request=request, text="RESOURCE_EXHAUSTED")
|
||||
client.post = AsyncMock(return_value=failed)
|
||||
|
||||
with pytest.raises(VertexAIError) as raised:
|
||||
await _async_call(endpoints, client)
|
||||
|
||||
assert raised.value.status_code == 429
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_waiters_resolve_on_their_own_when_the_first_request_fails(endpoints: ContextCachingEndpoints):
|
||||
client: Final = MagicMock(spec=AsyncHTTPHandler)
|
||||
|
||||
async def slow_list(**_):
|
||||
await asyncio.sleep(0.01)
|
||||
return _response({})
|
||||
|
||||
client.get = AsyncMock(side_effect=slow_list)
|
||||
client.post = AsyncMock(
|
||||
side_effect=[
|
||||
VertexAIError(status_code=500, message="create failed"),
|
||||
_response(_created("cachedContents/2", _expire_time(60))),
|
||||
]
|
||||
)
|
||||
|
||||
results: Final = await asyncio.gather(
|
||||
_async_call(endpoints, client), _async_call(endpoints, client), return_exceptions=True
|
||||
)
|
||||
|
||||
assert isinstance(results[0], VertexAIError)
|
||||
assert results[1][2] == "cachedContents/2"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("waiter_timeout", [0.05, httpx.Timeout(30.0, read=0.05)])
|
||||
@pytest.mark.asyncio
|
||||
async def test_waiter_resolves_on_its_own_once_its_timeout_passes(
|
||||
endpoints: ContextCachingEndpoints, waiter_timeout: float | httpx.Timeout
|
||||
):
|
||||
client: Final = MagicMock(spec=AsyncHTTPHandler)
|
||||
|
||||
async def slow_list(**_):
|
||||
await asyncio.sleep(0.5)
|
||||
return _response({})
|
||||
|
||||
client.get = AsyncMock(side_effect=slow_list)
|
||||
client.post = AsyncMock(return_value=_response(_created("cachedContents/1", _expire_time(60))))
|
||||
|
||||
first, waiter = await asyncio.gather(_async_call(endpoints, client), _async_call(endpoints, client, waiter_timeout))
|
||||
|
||||
assert (first[2], waiter[2]) == ("cachedContents/1", "cachedContents/1")
|
||||
assert client.get.await_count == 2
|
||||
|
||||
|
||||
def test_parse_expire_time_accepts_nanosecond_precision():
|
||||
assert parse_expire_time("2026-10-14T18:06:21.123456789Z") == pytest.approx(
|
||||
datetime(2026, 10, 14, 18, 6, 21, 123456, tzinfo=timezone.utc).timestamp()
|
||||
)
|
||||
assert (
|
||||
parse_expire_time("2026-10-14T18:06:21Z") == datetime(2026, 10, 14, 18, 6, 21, tzinfo=timezone.utc).timestamp()
|
||||
)
|
||||
assert parse_expire_time("not a time") is None
|
||||
assert parse_expire_time(None) is None
|
||||
|
||||
|
||||
def test_parse_expire_time_reads_a_missing_offset_as_utc():
|
||||
assert (
|
||||
parse_expire_time("2026-10-14T18:06:21") == datetime(2026, 10, 14, 18, 6, 21, tzinfo=timezone.utc).timestamp()
|
||||
)
|
||||
|
||||
|
||||
def test_cache_expiring_within_the_margin_is_not_remembered():
|
||||
known: Final = KnownCachedContents(clock=lambda: _NOW)
|
||||
known.remember(_SCOPE, "key", "cachedContents/soon", _expire_time(3))
|
||||
|
||||
assert known.get(_SCOPE, "key") is None
|
||||
|
||||
|
||||
def test_registry_is_bounded():
|
||||
known: Final = KnownCachedContents(clock=lambda: _NOW)
|
||||
with patch(f"{_REGISTRY_MODULE}.MAX_KNOWN_CACHED_CONTENTS", 3):
|
||||
for index in range(5):
|
||||
known.remember(_SCOPE, f"key-{index}", f"cachedContents/{index}", _expire_time(60))
|
||||
|
||||
assert known.get(_SCOPE, "key-0") is None
|
||||
assert known.get(_SCOPE, "key-4") == "cachedContents/4"
|
||||
|
||||
|
||||
def _generated(stream: bool) -> httpx.Response:
|
||||
body: Final = {
|
||||
"candidates": [{"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"}],
|
||||
"usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 1, "totalTokenCount": 11},
|
||||
}
|
||||
if stream:
|
||||
return httpx.Response(200, text=f"data: {json.dumps(body)}\n\n", headers={"content-type": "text/event-stream"})
|
||||
return httpx.Response(200, json=body)
|
||||
|
||||
|
||||
def _rejected(status_code: int, message: str) -> httpx.Response:
|
||||
return httpx.Response(status_code, json={"error": {"code": status_code, "message": message}})
|
||||
|
||||
|
||||
async def _complete(sync_mode: bool, stream: bool, **kwargs) -> str:
|
||||
request: Final = {"model": "gemini/gemini-3.1-pro", "messages": _messages(), "stream": stream, **kwargs}
|
||||
if sync_mode:
|
||||
response = litellm.completion(**request)
|
||||
if stream:
|
||||
return "".join(chunk.choices[0].delta.content or "" for chunk in response)
|
||||
return response.choices[0].message.content
|
||||
response = await litellm.acompletion(**request)
|
||||
if stream:
|
||||
return "".join([chunk.choices[0].delta.content or "" async for chunk in response])
|
||||
return response.choices[0].message.content
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.parametrize("stream", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_rejected_cache_is_replaced_when_a_request_param_cannot_be_deep_copied(
|
||||
respx_mock: respx.MockRouter,
|
||||
fake_provider_credentials: None,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sync_mode: bool,
|
||||
stream: bool,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
respx_mock.get(url__regex=r".*/cachedContents.*").mock(return_value=httpx.Response(200, json={}))
|
||||
respx_mock.post(url__regex=r".*/cachedContents$").mock(
|
||||
side_effect=[
|
||||
httpx.Response(200, json=_created("cachedContents/deleted", _FAR_FUTURE)),
|
||||
httpx.Response(200, json=_created("cachedContents/fresh", _FAR_FUTURE)),
|
||||
]
|
||||
)
|
||||
generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock(
|
||||
side_effect=[_generated(stream), _rejected(404, "CachedContent not found"), _generated(stream)]
|
||||
)
|
||||
lock: Final = threading.Lock()
|
||||
|
||||
assert await _complete(sync_mode, stream, lock_holder=lock) == "ok"
|
||||
assert await _complete(sync_mode, stream, lock_holder=lock) == "ok"
|
||||
|
||||
sent: Final = [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls]
|
||||
assert sent == ["cachedContents/deleted", "cachedContents/deleted", "cachedContents/fresh"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.parametrize("stream", [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
("status_code", "message", "replaced"),
|
||||
[
|
||||
(400, "Invalid resource state for cache content 4484391867902853120.", True),
|
||||
(403, "CachedContent not found (or permission denied)", True),
|
||||
(404, "CachedContent not found", True),
|
||||
(400, "Invalid value at 'generation_config.temperature'", False),
|
||||
(500, "CachedContent not found", False),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_remembered_cache_rejected_by_generate_content_is_replaced_within_the_same_call(
|
||||
respx_mock: respx.MockRouter,
|
||||
fake_provider_credentials: None,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sync_mode: bool,
|
||||
stream: bool,
|
||||
status_code: int,
|
||||
message: str,
|
||||
replaced: bool,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
list_route: Final = respx_mock.get(url__regex=r".*/cachedContents.*").mock(
|
||||
return_value=httpx.Response(200, json={})
|
||||
)
|
||||
respx_mock.post(url__regex=r".*/cachedContents$").mock(
|
||||
side_effect=[
|
||||
httpx.Response(200, json=_created("cachedContents/deleted", _FAR_FUTURE)),
|
||||
httpx.Response(200, json=_created("cachedContents/fresh", _FAR_FUTURE)),
|
||||
]
|
||||
)
|
||||
generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock(
|
||||
side_effect=[_generated(stream), _rejected(status_code, message), _generated(stream)]
|
||||
)
|
||||
logged: Final[list[dict]] = []
|
||||
|
||||
assert await _complete(sync_mode, stream) == "ok"
|
||||
if replaced:
|
||||
assert await _complete(sync_mode, stream, logger_fn=logged.append) == "ok"
|
||||
else:
|
||||
with pytest.raises(openai.APIStatusError):
|
||||
await _complete(sync_mode, stream, logger_fn=logged.append)
|
||||
|
||||
sent: Final = [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls]
|
||||
assert sent == ["cachedContents/deleted", "cachedContents/deleted"] + (["cachedContents/fresh"] if replaced else [])
|
||||
assert list_route.call_count == (2 if replaced else 1)
|
||||
assert "enerateContent" in logged[-1]["litellm_params"]["api_base"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.parametrize("stream", [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
("status_code", "message", "forgotten"),
|
||||
[
|
||||
(400, "Invalid value at 'generation_config.temperature'", True),
|
||||
(
|
||||
400,
|
||||
"Tool config, tools and system instruction should not be set in therequest when using cached content.",
|
||||
True,
|
||||
),
|
||||
(
|
||||
400,
|
||||
"CachedContent can not be used with GenerateContent request setting system_instruction, tools or tool_config.",
|
||||
True,
|
||||
),
|
||||
(404, "Requested entity was not found.", True),
|
||||
(429, "Resource has been exhausted (e.g. check quota).", False),
|
||||
(500, "Internal error encountered.", False),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_that_does_not_name_a_missing_cache_is_not_retried(
|
||||
respx_mock: respx.MockRouter,
|
||||
fake_provider_credentials: None,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sync_mode: bool,
|
||||
stream: bool,
|
||||
status_code: int,
|
||||
message: str,
|
||||
forgotten: bool,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
list_route: Final = respx_mock.get(url__regex=r".*/cachedContents.*").mock(
|
||||
return_value=httpx.Response(200, json={})
|
||||
)
|
||||
respx_mock.post(url__regex=r".*/cachedContents$").mock(
|
||||
side_effect=[
|
||||
httpx.Response(200, json=_created("cachedContents/first", _FAR_FUTURE)),
|
||||
httpx.Response(200, json=_created("cachedContents/second", _FAR_FUTURE)),
|
||||
]
|
||||
)
|
||||
generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock(
|
||||
side_effect=[_generated(stream), _rejected(status_code, message), _generated(stream)]
|
||||
)
|
||||
|
||||
assert await _complete(sync_mode, stream) == "ok"
|
||||
with pytest.raises(openai.APIStatusError):
|
||||
await _complete(sync_mode, stream)
|
||||
assert await _complete(sync_mode, stream) == "ok"
|
||||
|
||||
third: Final = "cachedContents/second" if forgotten else "cachedContents/first"
|
||||
assert [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] == [
|
||||
"cachedContents/first",
|
||||
"cachedContents/first",
|
||||
third,
|
||||
]
|
||||
assert list_route.call_count == (2 if forgotten else 1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.parametrize("stream", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_rejected_caller_supplied_cached_content_is_not_sent_again(
|
||||
respx_mock: respx.MockRouter,
|
||||
fake_provider_credentials: None,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sync_mode: bool,
|
||||
stream: bool,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
list_route: Final = respx_mock.get(url__regex=r".*/cachedContents.*").mock(
|
||||
return_value=httpx.Response(200, json={})
|
||||
)
|
||||
generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock(
|
||||
side_effect=[_rejected(404, "CachedContent not found"), _generated(stream)]
|
||||
)
|
||||
|
||||
with pytest.raises(openai.APIStatusError):
|
||||
await _complete(sync_mode, stream, cached_content="cachedContents/callers")
|
||||
|
||||
assert [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] == [
|
||||
"cachedContents/callers"
|
||||
]
|
||||
assert list_route.call_count == 0
|
||||
|
||||
|
||||
def _cut_off_stream() -> httpx.Response:
|
||||
body: Final = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"role": "model", "parts": [{"text": "ok"}]},
|
||||
"finishReason": "CONTINUATION",
|
||||
"continuationToken": "dG9rLTE=",
|
||||
}
|
||||
],
|
||||
"usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 1, "totalTokenCount": 11},
|
||||
}
|
||||
return httpx.Response(200, text=f"data: {json.dumps(body)}\n\n", headers={"content-type": "text/event-stream"})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_rejected_on_a_resume_hop_is_not_sent_again(
|
||||
respx_mock: respx.MockRouter,
|
||||
fake_provider_credentials: None,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sync_mode: bool,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
list_route: Final = respx_mock.get(url__regex=r".*/cachedContents.*").mock(
|
||||
return_value=httpx.Response(200, json={})
|
||||
)
|
||||
respx_mock.post(url__regex=r".*/cachedContents$").mock(
|
||||
side_effect=[
|
||||
httpx.Response(200, json=_created("cachedContents/deleted", _FAR_FUTURE)),
|
||||
httpx.Response(200, json=_created("cachedContents/fresh", _FAR_FUTURE)),
|
||||
]
|
||||
)
|
||||
generate_route: Final = respx_mock.post(url__regex=r".*:streamGenerateContent.*").mock(
|
||||
side_effect=[_cut_off_stream(), _rejected(404, "CachedContent not found"), _generated(True)]
|
||||
)
|
||||
|
||||
await _complete(sync_mode, True, max_tokens=100_000)
|
||||
await _complete(sync_mode, True, max_tokens=100_000)
|
||||
|
||||
sent: Final = [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls]
|
||||
assert sent == ["cachedContents/deleted", "cachedContents/deleted", "cachedContents/fresh"]
|
||||
assert list_route.call_count == 2
|
||||
Loading…
Add table
Reference in a new issue