fix(caching): read the cache key's model group from the router's metadata for the call (#48917)

* test(caching): audit #44542 response cache model group key

* fix(caching): key the response cache by the router's metadata bucket only

* fix(caching): read the router's model group before the caller's metadata

* fix(caching): read caching groups only beside the router's model group

A caller's own metadata.caching_groups on /v1/responses and /v1/messages could
still pick the cache key, so a string group let one model group fill another's
entry. The key now reads caching_groups only from the metadata that holds the
router's model_group.

* test(caching): type the planted caching groups parameter

* fix(caching): read the model group from the metadata that carries the call's deployment

A router call can carry a litellm_metadata it did not write, from the auto router's classifier embedding or an SDK caller, while the router writes its own group into metadata. The cache now picks the metadata source whose model_info matches the call's deployment, then falls back to the previous order.

* fix(caching): keep main's model group lookup for calls outside the router

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
mika-lt 2026-10-29 05:35:23 -07:00 • committed by GitHub
parent 41602e318d
commit ebf84f8649
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 326 additions and 14 deletions

View file

@ -455,17 +455,33 @@ class Cache:
3. Else use the `model` passed in kwargs, plus `model_id` when set (e.g. a SageMaker inference component)
"""
litellm_params: Final[dict] = kwargs.get("litellm_params", {}) or {}
metadata_sources: Final[tuple[dict, ...]] = (
kwargs.get("metadata") or {},
kwargs.get("litellm_metadata") or {},
litellm_params.get("metadata") or {},
litellm_params.get("litellm_metadata") or {},
metadata_sources: Final = cast( # cast-ok: each source is the untyped request's metadata dict or empty
"tuple[dict[str, object], dict[str, object], dict[str, object], dict[str, object]]",
(
kwargs.get("metadata") or {},
kwargs.get("litellm_metadata") or {},
litellm_params.get("metadata") or {},
litellm_params.get("litellm_metadata") or {},
),
)
model_group: Final[str | None] = next(
(source["model_group"] for source in metadata_sources if source.get("model_group")), None
metadata, litellm_metadata, params_metadata, params_litellm_metadata = metadata_sources
call_model_info: Final[object] = kwargs.get("model_info")
routed_source: Final = next(
(
source
for source in (litellm_metadata, metadata, params_litellm_metadata, params_metadata)
if call_model_info and source.get("model_group") and source.get("model_info") == call_model_info
),
None,
)
model_group: Final = (
str(routed_source["model_group"])
if routed_source is not None
else next((str(source["model_group"]) for source in metadata_sources if source.get("model_group")), None)
)
caching_group_sources: Final = (routed_source,) if routed_source is not None else metadata_sources
caching_group: Final = next(
(group for source in metadata_sources if (group := self._get_caching_group(source, model_group))), None
(group for source in caching_group_sources if (group := self._get_caching_group(source, model_group))), None
)
model_id_suffix: Final = (
f" model_id: {kwargs['model_id']}" if "model_id" in kwargs and kwargs["model_id"] is not None else ""

View file

@ -280,12 +280,14 @@ class Scenario:
self.gateway.rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == []
)
def model(self, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue) -> str:
name: Final = f"integration-{uuid.uuid4().hex}"
def model(
self, *, name: str | None = None, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue
) -> str:
model_name: Final = name if name is not None else f"integration-{uuid.uuid4().hex}"
created: Final = self.gateway.post(
"/model/new",
{
"model_name": name,
"model_name": model_name,
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "integration-provider-key",
@ -298,8 +300,8 @@ class Scenario:
identity: Final = string_value(object_value(created["model_info"])["id"])
self.cleanups.callback(self.delete_model, identity)
if self.gateway.workers > 1:
self._await_every_worker(name, identity)
return name
self._await_every_worker(model_name, identity)
return model_name
def _await_every_worker(self, name: str, identity: str) -> None:
from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS

View file

@ -0,0 +1,192 @@
"""The response cache keys an answer by the router's model group on every endpoint the router serves.
The router names the model group in ``metadata`` on ``/v1/chat/completions`` and in ``litellm_metadata`` on
``/v1/responses`` and ``/v1/messages``. Every cell gives each model group its own scripted upstream scenario on one
shared provider model, fills one group's entry, and reads which scenario the identical request reached next and which
answer the caller got: an answer served from the cache repeats the id of the request that filled it.
"""
from __future__ import annotations
import uuid
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Final, TypeAlias
import pytest
import yaml
from anthropic import Anthropic
from openai import OpenAI
from pydantic import JsonValue
from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, object_value
from tests.integration._support.process import OwnedState, graceful_stop_seconds, owned_proxy_process
from tests.integration._support.upstream import ScenarioHandle
from tests.integration.caching.response_cache_case import (
CHAT,
CHAT_MODEL,
MESSAGE,
MESSAGES_MODEL,
PROVIDER_KEY,
RESPONSE,
Upstream,
await_entry,
json_response,
prompt,
redis_store,
scripted,
)
Ask: TypeAlias = Callable[[Gateway, str, str, Mapping[str, JsonValue]], str]
def _root(gateway: Gateway) -> str:
return str(gateway.client.base_url).rstrip("/")
def _metadata(metadata: Mapping[str, JsonValue]) -> dict[str, JsonValue] | None:
return {"metadata": dict(metadata)} if metadata else None
def _chat(gateway: Gateway, model: str, text: str, metadata: Mapping[str, JsonValue]) -> str:
with OpenAI(base_url=f"{_root(gateway)}/v1", api_key=gateway.key, max_retries=0) as client:
return client.chat.completions.create(
model=model, messages=[{"role": "user", "content": text}], extra_body=_metadata(metadata)
).id
def _responses(gateway: Gateway, model: str, text: str, metadata: Mapping[str, JsonValue]) -> str:
with OpenAI(base_url=f"{_root(gateway)}/v1", api_key=gateway.key, max_retries=0) as client:
identity: Final = client.responses.create(model=model, input=text, extra_body=_metadata(metadata)).output[0].id
assert identity is not None
return identity
def _messages(gateway: Gateway, model: str, text: str, metadata: Mapping[str, JsonValue]) -> str:
with Anthropic(base_url=_root(gateway), api_key=gateway.key, max_retries=0) as client:
return client.messages.create(
model=model,
max_tokens=16,
messages=[{"role": "user", "content": text}],
extra_body=_metadata(metadata),
).id
@dataclass(frozen=True, slots=True)
class Endpoint:
provider_model: str
sibling_model: str
answer: Mapping[str, JsonValue]
ask: Ask
@dataclass(frozen=True, slots=True)
class Group:
name: str
upstream: ScenarioHandle
ENDPOINTS: Final = {
"chat_completions": Endpoint(CHAT_MODEL, "openai/gpt-5.6", CHAT, _chat),
"responses": Endpoint(CHAT_MODEL, "openai/gpt-5.6", RESPONSE, _responses),
"messages": Endpoint(MESSAGES_MODEL, "anthropic/claude-sonnet-5-5", MESSAGE, _messages),
}
def _group(scenario: Scenario, endpoint: Endpoint, name: str | None = None) -> Group:
handle: Final = scripted(scenario, json_response(endpoint.answer))
return Group(scenario.model(name=name, model=endpoint.provider_model, api_base=handle.api_base()), handle)
@pytest.mark.parametrize("endpoint", ENDPOINTS.values(), ids=ENDPOINTS.keys())
def test_identical_requests_to_two_groups_on_one_model_each_reach_their_own_deployment(
gateway: Gateway, endpoint: Endpoint
) -> None:
upstream: Final = Upstream(gateway.upstream_url)
store: Final = redis_store()
text: Final = prompt()
with gateway.scenario() as scenario:
first: Final = _group(scenario, endpoint)
second: Final = _group(scenario, endpoint)
filled: Final = endpoint.ask(gateway, first.name, text, {})
await_entry(store, filled)
answered: Final = endpoint.ask(gateway, second.name, text, {})
assert answered != filled, f"{second.name} was served the answer cached for {first.name}"
assert upstream.calls(second.upstream.scenario_id) == 1
assert endpoint.ask(gateway, first.name, text, {}) == filled
assert upstream.calls(first.upstream.scenario_id) == 1
PLANTS: Final[Mapping[str, Callable[[str], Mapping[str, JsonValue]]]] = {
"model_group": lambda target: {"model_group": target},
"caching_groups": lambda target: {"caching_groups": [target]},
}
@pytest.mark.parametrize("plant", PLANTS.values(), ids=PLANTS.keys())
@pytest.mark.parametrize("endpoint", ENDPOINTS.values(), ids=ENDPOINTS.keys())
def test_a_caller_naming_another_group_in_metadata_never_fills_that_groups_entry(
gateway: Gateway, endpoint: Endpoint, plant: Callable[[str], Mapping[str, JsonValue]]
) -> None:
upstream: Final = Upstream(gateway.upstream_url)
store: Final = redis_store()
text: Final = prompt()
with gateway.scenario() as scenario:
caller: Final = _group(scenario, endpoint)
target: Final = _group(scenario, endpoint, name=f"{caller.name}-premium")
planted: Final = endpoint.ask(gateway, caller.name, text, plant(target.name))
assert upstream.calls(caller.upstream.scenario_id) == 1
await_entry(store, planted)
answered: Final = endpoint.ask(gateway, target.name, text, {})
assert answered != planted, f"{target.name} was served the answer {caller.name} planted under its name"
assert upstream.calls(target.upstream.scenario_id) == 1
await_entry(store, answered)
assert endpoint.ask(gateway, target.name, text, {}) == answered
def _caching_groups_config(directory: Path, endpoint: Endpoint, first: Group, second: Group) -> Path:
rig: Final = JSON_OBJECT.validate_python(
yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text())
)
router_settings: Final = object_value(rig.get("router_settings") or {})
config: Final = {
**rig,
"model_list": [
{
"model_name": group.name,
"litellm_params": {"model": model, "api_key": PROVIDER_KEY, "api_base": group.upstream.api_base()},
}
for group, model in ((first, endpoint.provider_model), (second, endpoint.sibling_model))
],
"router_settings": {**router_settings, "caching_groups": [[first.name, second.name]]},
}
path: Final = directory / f"proxy_config-{uuid.uuid4().hex[:8]}.yaml"
path.write_text(yaml.safe_dump(config, sort_keys=False))
return path
def _configured_group(scenario: Scenario, endpoint: Endpoint) -> Group:
return Group(f"cache-group-{uuid.uuid4().hex[:8]}", scripted(scenario, json_response(endpoint.answer)))
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
@pytest.mark.parametrize("endpoint", ENDPOINTS.values(), ids=ENDPOINTS.keys())
def test_groups_in_one_caching_group_on_different_models_share_one_entry(
gateway: Gateway, endpoint: Endpoint, tmp_path: Path
) -> None:
upstream: Final = Upstream(gateway.upstream_url)
store: Final = redis_store()
text: Final = prompt()
with gateway.scenario() as scenario:
first: Final = _configured_group(scenario, endpoint)
second: Final = _configured_group(scenario, endpoint)
config: Final = _caching_groups_config(tmp_path, endpoint, first, second)
with owned_proxy_process(
gateway, tmp_path, {}, config=config, workers=1, state=OwnedState.SHARED_WITH_WORKER
) as owned:
filled: Final = endpoint.ask(owned.gateway, first.name, text, {})
await_entry(store, filled)
assert endpoint.ask(owned.gateway, second.name, text, {}) == filled
assert upstream.calls(second.upstream.scenario_id) == 0
assert upstream.calls(first.upstream.scenario_id) == 1

View file

@ -8,6 +8,7 @@ load_dotenv()
import asyncio
import hashlib
import random
from typing import Final
import pytest
@ -16,7 +17,7 @@ from litellm import aembedding, completion, embedding
from litellm.caching.caching import Cache
from unittest.mock import AsyncMock, patch, MagicMock
from litellm.caching.caching_handler import LLMCachingHandler, CachingHandlerResponse
from litellm.caching.caching_handler import LLMCachingHandler, CachingHandlerResponse, _PENDING_CACHE_WRITES
from litellm.caching.caching import LiteLLMCacheType
from litellm.types.utils import CallTypes
from litellm.types.rerank import RerankResponse
@ -278,6 +279,107 @@ def test_get_model_param_value_reads_model_group_from_litellm_metadata():
)
def test_get_model_param_value_prefers_the_router_group_over_a_caller_model_group():
cache: Final = Cache()
deployment: Final = {"id": "deployment-1"}
request: Final = {"model": "openai/gpt-5.6", "input": "search this text", "model_info": deployment}
router_metadata: Final = {"model_group": "caller", "model_info": deployment}
assert (
cache._get_model_param_value(
{**request, "metadata": {"model_group": "target"}, "litellm_metadata": router_metadata}
)
== "caller"
)
assert (
cache._get_model_param_value(
{
**request,
"litellm_params": {"metadata": {"model_group": "target"}, "litellm_metadata": router_metadata},
}
)
== "caller"
)
assert (
cache._get_model_param_value(
{
**request,
"metadata": {"model_group": "target", "model_info": deployment},
"litellm_metadata": router_metadata,
}
)
== "caller"
)
assert (
cache._get_model_param_value(
{
**request,
"metadata": router_metadata,
"litellm_metadata": {"model_group": "parent", "model_info": {"id": "parent-deployment"}},
}
)
== "caller"
)
def test_get_model_param_value_keeps_metadata_first_outside_the_router():
cache: Final = Cache()
assert (
cache._get_model_param_value(
{
"model": "openai/gpt-5.6",
"input": "search this text",
"metadata": {"model_group": "group-a"},
"litellm_metadata": {"model_group": "group-b"},
}
)
== "group-a"
)
@pytest.mark.asyncio
async def test_router_call_keeps_its_own_group_when_forwarded_litellm_metadata_names_another(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL))
router: Final = litellm.Router(
model_list=[
{
"model_name": group,
"litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-test", "mock_response": f"{group} answer"},
}
for group in ("group-a", "group-b")
]
)
messages: Final = [{"role": "user", "content": "which group answers this"}]
await router.acompletion(
model="group-a",
messages=messages,
caching=True,
litellm_metadata={"model_group": "group-b", "model_info": {"id": "parent-deployment"}},
)
await asyncio.gather(*_PENDING_CACHE_WRITES)
reply: Final = await router.acompletion(model="group-b", messages=messages, caching=True)
assert reply.choices[0].message.content == "group-b answer"
@pytest.mark.parametrize("planted", [["sales-premium"], [("sales", "sales-premium")]], ids=["string", "tuple"])
def test_get_model_param_value_reads_caching_groups_only_beside_the_router_group(planted: list[str | tuple[str, str]]):
cache: Final = Cache()
request: Final = {
"model": "openai/gpt-5.6",
"input": "search this text",
"model_info": {"id": "sales-deployment"},
"litellm_metadata": {"model_group": "sales", "caching_groups": None, "model_info": {"id": "sales-deployment"}},
"metadata": {"caching_groups": planted},
}
assert cache._get_model_param_value(request) == "sales"
def test_preset_cache_key():
"""
Test that the preset cache key is used if it is set in kwargs["litellm_params"]