scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

fix(vertex_ai): retry only errors that name the cache, and log the generateContent request after a rebuild

Browse source
This commit is contained in:
mateo-berri 2026-10-20 00:53:23 -07:00
parent a7ba6c9ef0
commit 6a5fea1ad8
4 changed files with 106 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/llms/vertex_ai/context_caching/known_cached_contents.py
View file

@ -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}

14
litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py
View file

@ -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)

71
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -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),

61
tests/unit/llms/vertex_ai/context_caching/test_known_cached_contents.py
View file

@ -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):