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

This commit is contained in:
mateo-berri 2026-10-29 01:20:48 -07:00
parent f92ab299e8
commit 35500573ec
2 changed files with 62 additions and 34 deletions

View file

@ -456,30 +456,37 @@ class Cache:
"""
litellm_params: Final[dict] = kwargs.get("litellm_params", {}) or {}
metadata_sources: Final = cast( # cast-ok: each source is the untyped request's metadata dict or empty
"tuple[dict[str, object], ...]",
"tuple[dict[str, object], dict[str, object], dict[str, object], dict[str, object]]",
(
kwargs.get("litellm_metadata") or {},
kwargs.get("metadata") or {},
litellm_params.get("litellm_metadata") or {},
kwargs.get("litellm_metadata") or {},
litellm_params.get("metadata") or {},
litellm_params.get("litellm_metadata") or {},
),
)
metadata, litellm_metadata, params_metadata, params_litellm_metadata = metadata_sources
call_model_info: Final[object] = kwargs.get("model_info")
model_group_source: Final = next(
routed_source: Final = next(
(
source
for source in metadata_sources
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,
) or next((source for source in metadata_sources if source.get("model_group")), None)
if model_group_source is None:
model_id_suffix: Final = (
f" model_id: {kwargs['model_id']}" if "model_id" in kwargs and kwargs["model_id"] is not None else ""
)
return f"{kwargs['model']}{model_id_suffix}"
model_group: Final = str(model_group_source["model_group"])
return self._get_caching_group(model_group_source, model_group) or model_group
)
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 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 ""
)
return caching_group or model_group or f"{kwargs['model']}{model_id_suffix}"
def _get_caching_group(self, metadata: dict, model_group: str | None) -> str | None:
caching_groups: Final[list | None] = metadata.get("caching_groups", [])

View file

@ -281,11 +281,13 @@ 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()
request: Final = {"model": "openai/gpt-5.6", "input": "search this text"}
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": {"model_group": "caller"}}
{**request, "metadata": {"model_group": "target"}, "litellm_metadata": router_metadata}
)
== "caller"
)
@ -293,32 +295,47 @@ def test_get_model_param_value_prefers_the_router_group_over_a_caller_model_grou
cache._get_model_param_value(
{
**request,
"litellm_params": {
"metadata": {"model_group": "target"},
"litellm_metadata": {"model_group": "caller"},
},
"litellm_params": {"metadata": {"model_group": "target"}, "litellm_metadata": router_metadata},
}
)
== "caller"
)
assert (
cache._get_model_param_value(
{**request, "metadata": {"model_group": "group-a"}, "litellm_metadata": {"attempted_retries": 0}}
{
**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"
)
deployment: Final = {"id": "deployment-1"}
assert (
cache._get_model_param_value(
{
**request,
"model_info": deployment,
"metadata": {"model_group": "target", "model_info": deployment},
"litellm_metadata": {"model_group": "caller", "model_info": deployment},
}
)
== "caller"
)
@pytest.mark.asyncio
@ -338,7 +355,10 @@ async def test_router_call_keeps_its_own_group_when_forwarded_litellm_metadata_n
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="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)
@ -352,7 +372,8 @@ def test_get_model_param_value_reads_caching_groups_only_beside_the_router_group
request: Final = {
"model": "openai/gpt-5.6",
"input": "search this text",
"litellm_metadata": {"model_group": "sales", "caching_groups": None},
"model_info": {"id": "sales-deployment"},
"litellm_metadata": {"model_group": "sales", "caching_groups": None, "model_info": {"id": "sales-deployment"}},
"metadata": {"caching_groups": planted},
}