scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-27 06:13:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 72

fix(vertex_ai): send once more with a fresh cache when generateContent rejects a remembered one

Browse source
This commit is contained in:
mateo-berri 2026-10-19 23:45:58 -07:00
parent c8706c8b03
commit 91acf19985
3 changed files with 227 additions and 65 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

74
litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py Normal file
View file

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

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

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

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

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