scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

fix(vertex_ai): copy request params per key before a cache retry so an uncopyable param no longer fails the call

Browse source 
cache_rebuild deep-copied optional_params as a whole, so an SDK call carrying a param deepcopy cannot copy (an unknown kwarg holding a threading.Lock) raised APIConnectionError on any cache-marked request. Each key is now copied on its own, and a key that cannot be copied keeps its reference, as safe_deep_copy does elsewhere
This commit is contained in:
mateo-berri 2026-10-20 03:52:43 -07:00
parent d9b0e221f5
commit 6d85aae12e
2 changed files with 33 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -5,6 +5,7 @@ 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
@ -22,7 +23,7 @@ def cache_rebuild(
"""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, **deepcopy({"messages": messages, "optional_params": optional_params}))
return partial(build, messages=deepcopy(messages), optional_params=independent_snapshot(optional_params))
def _status_and_message(error: httpx.HTTPStatusError | BaseLLMException) -> tuple[int, str]:

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

@ -1,5 +1,6 @@
import asyncio
import json
import threading
from datetime import datetime, timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
@ -450,6 +451,36 @@ async def _complete(sync_mode: bool, stream: bool, **kwargs) -> str:
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(