mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(proxy): keep a request's cache opt-in when its key has cache metadata (#50999)
* fix(proxy): merge a key's cache controls into the request's cache settings instead of replacing them * fix(proxy): keep a key with no-cache out of opt-in caching A key whose cache metadata sets no-cache now drops the caller's use-cache, so on a default_off proxy its responses stay out of the shared cache as before the merge change. default_on proxies ignore use-cache, so they are unchanged * fix(proxy): keep base handling of non-dict key cache metadata Key cache metadata that is not a dict, such as false or a string, wipes the request's cache settings as before the merge change, so those keys stay out of opt-in caching * fix(proxy): keep a key with use-cache false out of opt-in caching * fix(proxy): only keep request cache settings on an allowed opt-in --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
quietjules
GitHub
parent
478d9bf5d2
commit
63d39b41fd
3 changed files with 190 additions and 8 deletions
|
|
@ -35,7 +35,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
|
|||
tpd_limit: int | None = None
|
||||
budget_duration: str | None = None
|
||||
budget_reset_at: datetime | None = None
|
||||
allowed_cache_controls: list | None = []
|
||||
allowed_cache_controls: list[object] | None = []
|
||||
allowed_routes: list | None = []
|
||||
key_type: str | None = None
|
||||
permissions: dict = {}
|
||||
|
|
|
|||
|
|
@ -2217,15 +2217,41 @@ class LiteLLMProxyRequestSetup:
|
|||
return data
|
||||
|
||||
@staticmethod
|
||||
def add_key_level_controls(key_metadata: dict | None, data: dict, _metadata_variable_name: str):
|
||||
def _request_cache_kept_for_key(
|
||||
request_cache: Mapping[str, object],
|
||||
key_cache: Mapping[str, object],
|
||||
allowed_cache_controls: Sequence[object] | None,
|
||||
) -> Mapping[str, object]:
|
||||
key_opts_out: Final = (
|
||||
key_cache.get("no-cache", False) is not False or key_cache.get("use-cache", True) is not True
|
||||
)
|
||||
request_is_allowed: Final = not allowed_cache_controls or all(
|
||||
control in allowed_cache_controls for control in request_cache
|
||||
)
|
||||
if key_opts_out or request_cache.get("use-cache") is not True or not request_is_allowed:
|
||||
return MappingProxyType({})
|
||||
return {k: v for k, v in request_cache.items() if not (k == "s-max-age" and "s-maxage" in key_cache)}
|
||||
|
||||
@staticmethod
|
||||
def add_key_level_controls(
|
||||
key_metadata: dict | None,
|
||||
data: dict,
|
||||
_metadata_variable_name: str,
|
||||
allowed_cache_controls: Sequence[object] | None = None,
|
||||
):
|
||||
if key_metadata is None:
|
||||
return data
|
||||
if "cache" in key_metadata:
|
||||
data["cache"] = {}
|
||||
if isinstance(key_metadata["cache"], dict):
|
||||
for k, v in key_metadata["cache"].items():
|
||||
if k in SupportedCacheControls:
|
||||
data["cache"][k] = v
|
||||
if "cache" in key_metadata and not isinstance(key_metadata["cache"], dict):
|
||||
data["cache"] = {} # rebind-ok: data is an out-param
|
||||
elif "cache" in key_metadata:
|
||||
request_cache: Final = _DECODED_JSON.validate_python(data["cache"]) if "cache" in data else None
|
||||
key_cache: Final = _request_fields(_DECODED_JSON.validate_python(key_metadata["cache"]))
|
||||
data["cache"] = { # rebind-ok: data is an out-param
|
||||
**LiteLLMProxyRequestSetup._request_cache_kept_for_key(
|
||||
_request_fields(request_cache), key_cache, allowed_cache_controls
|
||||
),
|
||||
**{k: v for k, v in key_cache.items() if k in SupportedCacheControls},
|
||||
}
|
||||
|
||||
## KEY-LEVEL SPEND LOGS / TAGS
|
||||
if "tags" in key_metadata and key_metadata["tags"] is not None:
|
||||
|
|
@ -2967,6 +2993,7 @@ async def add_litellm_data_to_request(
|
|||
key_metadata=key_metadata,
|
||||
data=data,
|
||||
_metadata_variable_name=_metadata_variable_name,
|
||||
allowed_cache_controls=user_api_key_dict.allowed_cache_controls,
|
||||
)
|
||||
data[_metadata_variable_name] = LiteLLMProxyRequestSetup.add_team_and_project_level_controls(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from pydantic import ValidationError as PydanticValidationError
|
|||
from starlette.datastructures import Headers
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import Cache, CacheMode
|
||||
from litellm.constants import (
|
||||
CALLER_SESSION_ID_METADATA_KEY,
|
||||
COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY,
|
||||
|
|
@ -64,6 +65,7 @@ from litellm.proxy.litellm_pre_call_utils import (
|
|||
resolve_tenant_otel_destinations,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
||||
from litellm.types.caching import LiteLLMCacheType
|
||||
from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES
|
||||
from litellm.types.utils import CredentialItem
|
||||
from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs
|
||||
|
|
@ -8379,6 +8381,159 @@ async def test_key_level_callback_vars_survive_the_strip():
|
|||
assert updated["dd_site"] == "us5.datadoghq.com"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"key_cache, request_cache, expected_cache, uses_cache_when_default_off",
|
||||
[
|
||||
pytest.param(
|
||||
{"ttl": 600},
|
||||
{"use-cache": True, "namespace": "team-a"},
|
||||
{"use-cache": True, "namespace": "team-a", "ttl": 600},
|
||||
True,
|
||||
id="key-control-added-to-request-opt-in",
|
||||
),
|
||||
pytest.param(
|
||||
{"no-store": True},
|
||||
{"use-cache": True, "no-store": False},
|
||||
{"use-cache": True, "no-store": True},
|
||||
True,
|
||||
id="key-control-wins-over-request",
|
||||
),
|
||||
pytest.param(
|
||||
{"use-cache": True, "namespace": "key-ns"},
|
||||
None,
|
||||
{},
|
||||
False,
|
||||
id="unsupported-key-controls-ignored",
|
||||
),
|
||||
pytest.param(None, {"use-cache": True}, {}, False, id="non-dict-key-cache-still-wipes-request"),
|
||||
pytest.param(
|
||||
{"no-cache": True},
|
||||
{"use-cache": True, "ttl": 60},
|
||||
{"no-cache": True},
|
||||
False,
|
||||
id="key-no-cache-keeps-the-key-out-of-opt-in-caching",
|
||||
),
|
||||
pytest.param(
|
||||
{"use-cache": False},
|
||||
{"use-cache": True, "ttl": 60},
|
||||
{},
|
||||
False,
|
||||
id="key-use-cache-false-keeps-the-key-out-of-opt-in-caching",
|
||||
),
|
||||
pytest.param(
|
||||
{"use-cache": True},
|
||||
{"use-cache": True},
|
||||
{"use-cache": True},
|
||||
True,
|
||||
id="key-use-cache-true-leaves-the-request-opt-in",
|
||||
),
|
||||
pytest.param(
|
||||
{"no-cache": "true"},
|
||||
{"use-cache": True},
|
||||
{"no-cache": "true"},
|
||||
False,
|
||||
id="key-non-boolean-no-cache-keeps-the-key-out-of-opt-in-caching",
|
||||
),
|
||||
pytest.param(
|
||||
{"no-cache": False},
|
||||
{"use-cache": True},
|
||||
{"use-cache": True, "no-cache": False},
|
||||
True,
|
||||
id="key-no-cache-false-leaves-the-request-opt-in",
|
||||
),
|
||||
pytest.param(
|
||||
{"ttl": 600},
|
||||
{"no-cache": True, "namespace": "team-a"},
|
||||
{"ttl": 600},
|
||||
False,
|
||||
id="request-without-opt-in-is-still-replaced",
|
||||
),
|
||||
pytest.param(
|
||||
{"ttl": 600},
|
||||
{"use-cache": "true", "no-store": True},
|
||||
{"ttl": 600},
|
||||
False,
|
||||
id="request-non-boolean-use-cache-is-still-replaced",
|
||||
),
|
||||
pytest.param(
|
||||
{"s-maxage": 1},
|
||||
{"use-cache": True, "s-max-age": 100000},
|
||||
{"use-cache": True, "s-maxage": 1},
|
||||
True,
|
||||
id="key-s-maxage-overrides-the-request-s-max-age-alias",
|
||||
),
|
||||
pytest.param(
|
||||
{"ttl": 600},
|
||||
{"use-cache": True, "s-max-age": 5},
|
||||
{"use-cache": True, "s-max-age": 5, "ttl": 600},
|
||||
True,
|
||||
id="request-s-max-age-kept-when-the-key-has-no-s-maxage",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_key_cache_controls_merge_into_request_cache_controls(
|
||||
key_cache: dict[str, object] | None,
|
||||
request_cache: dict[str, object] | None,
|
||||
expected_cache: dict[str, object],
|
||||
uses_cache_when_default_off: bool,
|
||||
):
|
||||
key: Final = UserAPIKeyAuth(api_key="hashed-key", metadata={"cache": key_cache})
|
||||
data: Final = {
|
||||
"model": "gpt-4.1-mini",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
**({} if request_cache is None else {"cache": request_cache}),
|
||||
}
|
||||
|
||||
updated: Final = await add_litellm_data_to_request(
|
||||
data=data,
|
||||
request=_callback_credential_request_mock(),
|
||||
user_api_key_dict=key,
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={},
|
||||
version="test-version",
|
||||
)
|
||||
|
||||
assert updated["cache"] == expected_cache
|
||||
opt_in_cache: Final = Cache(type=LiteLLMCacheType.LOCAL, mode=CacheMode.default_off)
|
||||
assert opt_in_cache.should_use_cache(**updated) is uses_cache_when_default_off
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"allowed_cache_controls, request_cache, expected_cache",
|
||||
[
|
||||
pytest.param(["no-store"], {"use-cache": True}, {"ttl": 600}, id="opt-in-not-allowed-is-replaced"),
|
||||
pytest.param(
|
||||
["use-cache"], {"use-cache": True, "no-store": True}, {"ttl": 600}, id="any-disallowed-control-is-replaced"
|
||||
),
|
||||
pytest.param(["use-cache"], {"use-cache": True}, {"use-cache": True, "ttl": 600}, id="allowed-opt-in-is-kept"),
|
||||
pytest.param([], {"use-cache": True}, {"use-cache": True, "ttl": 600}, id="empty-list-allows-the-opt-in"),
|
||||
],
|
||||
)
|
||||
async def test_key_cache_metadata_only_keeps_request_controls_the_key_allows(
|
||||
allowed_cache_controls: list[str],
|
||||
request_cache: dict[str, object],
|
||||
expected_cache: dict[str, object],
|
||||
):
|
||||
key: Final = UserAPIKeyAuth(
|
||||
api_key="hashed-key",
|
||||
metadata={"cache": {"ttl": 600}},
|
||||
allowed_cache_controls=allowed_cache_controls,
|
||||
)
|
||||
|
||||
updated: Final = await add_litellm_data_to_request(
|
||||
data={"model": "gpt-4.1-mini", "messages": [{"role": "user", "content": "hello"}], "cache": request_cache},
|
||||
request=_callback_credential_request_mock(),
|
||||
user_api_key_dict=key,
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={},
|
||||
version="test-version",
|
||||
)
|
||||
|
||||
assert updated["cache"] == expected_cache
|
||||
|
||||
|
||||
class TestPromotedTraceControlFields:
|
||||
"""LIT-5137: caller metadata trace fields must reach litellm_metadata."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue