From 91acf1998516cb86f494e38cf739347a2e5fd5fb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 23:45:58 -0700 Subject: [PATCH] fix(vertex_ai): send once more with a fresh cache when generateContent rejects a remembered one --- .../context_caching/rejected_cache_retry.py | 74 ++++++++++ .../vertex_and_google_ai_studio_gemini.py | 136 ++++++++++++------ .../test_known_cached_contents.py | 82 ++++++++--- 3 files changed, 227 insertions(+), 65 deletions(-) create mode 100644 litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py diff --git a/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py b/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py new file mode 100644 index 00000000000..61c85f5fcdf --- /dev/null +++ b/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py @@ -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.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.llms.vertex_ai import RequestBody +from litellm.utils import is_cached_message + +from .known_cached_contents import CACHE_REJECTED_STATUS_CODES + +Built: Final = TypeVar("Built") +Sent: Final = TypeVar("Sent") + + +def cache_rebuild(build: Callable[..., Built], transform_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, transform_params["messages"])): + return None + return partial( + build, + **{ + **transform_params, + **deepcopy( + {"messages": transform_params["messages"], "optional_params": transform_params["optional_params"]} + ), + }, + ) + + +def _status_code(error: httpx.HTTPStatusError | BaseLLMException) -> int: + match error: + case httpx.HTTPStatusError(): + return error.response.status_code + case BaseLLMException(): + return error.status_code + + +def _cache_rejected(body: RequestBody, error: httpx.HTTPStatusError | BaseLLMException) -> bool: + return body.get("cachedContent") is not None and _status_code(error) in CACHE_REJECTED_STATUS_CODES + + +async def post_with_fresh_cache_async( + body: RequestBody, + rebuild: Callable[[], Awaitable[RequestBody]] | None, + post: Callable[[RequestBody], Awaitable[Sent]], +) -> 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() + 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], +) -> 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() + if fresh_body.get("cachedContent") == body.get("cachedContent"): + raise + return fresh_body, post(fresh_body) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index f30e8c69324..ca12ee54a80 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -97,6 +97,11 @@ from ..common_utils import ( 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 ( @@ -2787,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", } @@ -2989,14 +2994,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, @@ -3039,6 +3044,15 @@ class VertexLLM(VertexBase): litellm_params=litellm_params, ) + rebuild: Final = cache_rebuild( + partial( + async_transform_request_body, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_auth_header=auth_header, + ), + data, + ) request_body: Final = await async_transform_request_body( **data, vertex_project=vertex_project, @@ -3057,29 +3071,34 @@ class VertexLLM(VertexBase): }, ) - request_body_str: Final = json.dumps(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, + ) + + 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))[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, @@ -3089,7 +3108,7 @@ class VertexLLM(VertexBase): async def async_completion( self, model: str, - messages: list, + messages: list[AllMessageValues], model_response: ModelResponse, print_verbose: Callable, data: dict, @@ -3099,7 +3118,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, @@ -3142,6 +3161,15 @@ class VertexLLM(VertexBase): litellm_params=litellm_params, ) + rebuild: Final = cache_rebuild( + partial( + async_transform_request_body, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_auth_header=auth_header, + ), + data, + ) request_body: Final = await async_transform_request_body( **data, vertex_project=vertex_project, @@ -3195,14 +3223,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) 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, ) @@ -3327,6 +3355,15 @@ class VertexLLM(VertexBase): ) ## TRANSFORMATION ## + rebuild: Final = cache_rebuild( + partial( + sync_transform_request_body, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_auth_header=auth_header, + ), + transform_request_params, + ) data: Final = sync_transform_request_body( **transform_request_params, vertex_project=vertex_project, @@ -3351,22 +3388,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)[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, @@ -3406,14 +3448,14 @@ class VertexLLM(VertexBase): ) return response - first_response: Final = post(data) + sent_body, first_response = post_with_fresh_cache_sync(data, rebuild, post) 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, ) diff --git a/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py b/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py index 8f9c4ff05d9..4c66615b705 100644 --- a/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py +++ b/tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py @@ -322,51 +322,97 @@ def test_registry_is_bounded(): assert known_cached_contents.get(_SCOPE, "key-4") == "cachedContents/4" -async def _complete(sync_mode: bool, stream: bool) -> None: - if sync_mode: - response = litellm.completion(model="gemini/gemini-3.1-pro", messages=_messages(), stream=stream) - for _chunk in response if stream else (): - pass - return - response = await litellm.acompletion(model="gemini/gemini-3.1-pro", messages=_messages(), stream=stream) +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: - async for _chunk in response: - pass + 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) -> httpx.Response: + return httpx.Response(status_code, json={"error": {"code": status_code, "message": "rejected"}}) + + +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.parametrize( - ("status_code", "relisted"), + ("status_code", "replaced"), [(400, True), (403, True), (404, True), (429, False), (500, False)], ) @pytest.mark.asyncio -async def test_cache_rejected_by_generate_content_is_looked_up_again( +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, - relisted: bool, + 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( - return_value=httpx.Response(200, json=_created("cachedContents/deleted", _expire_time(300))) + side_effect=[ + httpx.Response(200, json=_created("cachedContents/deleted", _expire_time(300))), + httpx.Response(200, json=_created("cachedContents/fresh", _expire_time(300))), + ] ) generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock( - return_value=httpx.Response(status_code, json={"error": {"code": status_code, "message": "rejected"}}) + side_effect=[_generated(stream), _rejected(status_code), _generated(stream)] ) - for _ in range(2): + assert await _complete(sync_mode, stream) == "ok" + if replaced: + assert await _complete(sync_mode, stream) == "ok" + else: with pytest.raises(openai.APIStatusError): await _complete(sync_mode, stream) + 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) + + +@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), _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/deleted", - "cachedContents/deleted", + "cachedContents/callers" ] - assert list_route.call_count == (2 if relisted else 1) + assert list_route.call_count == 0