mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-07 07:58:09 +00:00
fix(caching): keep main's model group lookup for calls outside the router
This commit is contained in:
parent
f92ab299e8
commit
35500573ec
2 changed files with 62 additions and 34 deletions
|
|
@ -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", [])
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue