diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 3778dee46c2..5e13c656496 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -455,19 +455,30 @@ 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("litellm_metadata") or {}, - kwargs.get("metadata") or {}, - litellm_params.get("litellm_metadata") or {}, - litellm_params.get("metadata") or {}, + metadata_sources: Final = cast( # cast-ok: each source is the untyped request's metadata dict or empty + "tuple[dict[str, object], ...]", + ( + kwargs.get("litellm_metadata") or {}, + kwargs.get("metadata") or {}, + litellm_params.get("litellm_metadata") or {}, + litellm_params.get("metadata") or {}, + ), ) - model_group_source: Final = next((source for source in metadata_sources if source.get("model_group")), None) + call_model_info: Final[object] = kwargs.get("model_info") + model_group_source: Final = next( + ( + source + for source in metadata_sources + 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"] + model_group: Final = str(model_group_source["model_group"]) return self._get_caching_group(model_group_source, model_group) or model_group def _get_caching_group(self, metadata: dict, model_group: str | None) -> str | None: diff --git a/tests/unit/caching/test_unit_test_caching.py b/tests/unit/caching/test_unit_test_caching.py index 7c302d0578d..28b959c427a 100644 --- a/tests/unit/caching/test_unit_test_caching.py +++ b/tests/unit/caching/test_unit_test_caching.py @@ -17,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 @@ -307,6 +307,43 @@ def test_get_model_param_value_prefers_the_router_group_over_a_caller_model_grou ) == "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 +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"} + ) + 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"])