From 6a5fea1ad8a960c65840e469059e1dc37c0ce478 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:53:23 -0700 Subject: [PATCH] fix(vertex_ai): retry only errors that name the cache, and log the generateContent request after a rebuild --- .../context_caching/known_cached_contents.py | 11 ++- .../context_caching/rejected_cache_retry.py | 14 ++-- .../vertex_and_google_ai_studio_gemini.py | 71 ++++++++++--------- .../test_known_cached_contents.py | 61 +++++++++++++--- 4 files changed, 106 insertions(+), 51 deletions(-) diff --git a/litellm/llms/vertex_ai/context_caching/known_cached_contents.py b/litellm/llms/vertex_ai/context_caching/known_cached_contents.py index 9491b464b61..eefff5d5beb 100644 --- a/litellm/llms/vertex_ai/context_caching/known_cached_contents.py +++ b/litellm/llms/vertex_ai/context_caching/known_cached_contents.py @@ -9,11 +9,16 @@ 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}) +_CACHE_REJECTED_STATUS_CODES: Final = frozenset({400, 403, 404}) +_NAMES_CACHED_CONTENT: Final = re.compile(r"cached?\s*content", 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 + + 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: @@ -91,8 +96,8 @@ class KnownCachedContents: 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: + def forget_rejected(self, cached_content: str | None, status_code: int, message: str) -> None: + if cached_content is None or not names_rejected_cache(status_code, message): return with self._lock: self._entries = {key: entry for key, entry in self._entries.items() if entry[0] != cached_content} diff --git a/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py b/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py index 81e8b62d9f0..5bb13684470 100644 --- a/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py +++ b/litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py @@ -10,7 +10,7 @@ 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 CACHE_REJECTED_STATUS_CODES +from .known_cached_contents import names_rejected_cache Built: Final = TypeVar("Built") Sent: Final = TypeVar("Sent") @@ -25,22 +25,23 @@ def cache_rebuild( return partial(build, **deepcopy({"messages": messages, "optional_params": optional_params})) -def _status_code(error: httpx.HTTPStatusError | BaseLLMException) -> int: +def _status_and_message(error: httpx.HTTPStatusError | BaseLLMException) -> tuple[int, str]: match error: case httpx.HTTPStatusError(): - return error.response.status_code + return error.response.status_code, error.response.text case BaseLLMException(): - return error.status_code + return error.status_code, error.message 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 + 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) @@ -48,6 +49,7 @@ async def post_with_fresh_cache_async( 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) @@ -57,6 +59,7 @@ 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) @@ -64,6 +67,7 @@ def post_with_fresh_cache_sync( 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) 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 3e627f0c091..1e5baac22cf 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 @@ -2806,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 @@ -2837,8 +2851,10 @@ async def make_call( ) response.raise_for_status() except httpx.HTTPStatusError as e: - known_cached_contents.forget_rejected(request_body.get("cachedContent"), e.response.status_code) exception_string: Final = str(await e.response.aread()) + known_cached_contents.forget_rejected( + request_body.get("cachedContent"), e.response.status_code, exception_string + ) raise VertexAIError( status_code=e.response.status_code, message=VertexGeminiConfig().translate_exception_str(exception_string), @@ -2908,7 +2924,9 @@ def make_sync_call( timeout=timeout, ) except httpx.HTTPStatusError as e: - known_cached_contents.forget_rejected(request_body.get("cachedContent"), e.response.status_code) + known_cached_contents.forget_rejected( + request_body.get("cachedContent"), e.response.status_code, e.response.text + ) raise if response.status_code != 200 and response.status_code != 201: @@ -3063,15 +3081,8 @@ 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, @@ -3096,7 +3107,7 @@ class VertexLLM(VertexBase): 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] + return (await post_with_fresh_cache_async(request_body, rebuild, post, log_request))[1] streaming_response: Final = CustomStreamWrapper( completion_stream=None, @@ -3189,15 +3200,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 @@ -3213,7 +3217,9 @@ class VertexLLM(VertexBase): ) response.raise_for_status() except httpx.HTTPStatusError as err: - known_cached_contents.forget_rejected(body.get("cachedContent"), err.response.status_code) + known_cached_contents.forget_rejected( + body.get("cachedContent"), err.response.status_code, err.response.text + ) raise VertexAIError( status_code=err.response.status_code, message=err.response.text, @@ -3227,7 +3233,7 @@ class VertexLLM(VertexBase): ) return response - sent_body, first_response = await post_with_fresh_cache_async(request_body, rebuild, post) + 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), sent_body, post), @@ -3378,15 +3384,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 @@ -3410,7 +3409,7 @@ class VertexLLM(VertexBase): 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] + return post_with_fresh_cache_sync(data, rebuild, post, log_request)[1] streaming_response: Final = CustomStreamWrapper( completion_stream=None, @@ -3440,7 +3439,9 @@ class VertexLLM(VertexBase): ) response.raise_for_status() except httpx.HTTPStatusError as err: - known_cached_contents.forget_rejected(body.get("cachedContent"), err.response.status_code) + known_cached_contents.forget_rejected( + body.get("cachedContent"), err.response.status_code, err.response.text + ) raise VertexAIError( status_code=err.response.status_code, message=err.response.text, @@ -3454,7 +3455,7 @@ class VertexLLM(VertexBase): ) return response - sent_body, first_response = post_with_fresh_cache_sync(data, rebuild, post) + 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), sent_body, post), 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 e3668aeebf3..14216c443b5 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 @@ -347,8 +347,8 @@ def _generated(stream: bool) -> httpx.Response: 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"}}) +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: @@ -367,8 +367,14 @@ async def _complete(sync_mode: bool, stream: bool, **kwargs) -> str: @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.parametrize("stream", [True, False]) @pytest.mark.parametrize( - ("status_code", "replaced"), - [(400, True), (403, True), (404, True), (429, False), (500, False)], + ("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( @@ -378,6 +384,7 @@ async def test_remembered_cache_rejected_by_generate_content_is_replaced_within_ sync_mode: bool, stream: bool, status_code: int, + message: str, replaced: bool, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) @@ -391,19 +398,57 @@ async def test_remembered_cache_rejected_by_generate_content_is_replaced_within_ ] ) generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock( - side_effect=[_generated(stream), _rejected(status_code), _generated(stream)] + 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) == "ok" + assert await _complete(sync_mode, stream, logger_fn=logged.append) == "ok" else: with pytest.raises(openai.APIStatusError): - await _complete(sync_mode, stream) + 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.asyncio +async def test_error_that_does_not_name_the_cache_keeps_the_remembered_cache( + 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={}) + ) + respx_mock.post(url__regex=r".*/cachedContents$").mock( + return_value=httpx.Response(200, json=_created("cachedContents/kept", _FAR_FUTURE)) + ) + generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock( + side_effect=[ + _generated(stream), + _rejected(400, "Invalid value at 'generation_config.temperature'"), + _generated(stream), + ] + ) + + assert await _complete(sync_mode, stream) == "ok" + with pytest.raises(openai.BadRequestError): + await _complete(sync_mode, stream) + assert await _complete(sync_mode, stream) == "ok" + + assert [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] == [ + "cachedContents/kept" + ] * 3 + assert list_route.call_count == 1 @pytest.mark.parametrize("sync_mode", [True, False]) @@ -421,7 +466,7 @@ async def test_rejected_caller_supplied_cached_content_is_not_sent_again( 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)] + side_effect=[_rejected(404, "CachedContent not found"), _generated(stream)] ) with pytest.raises(openai.APIStatusError):