diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 5f5b8210096..1ced13445de 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -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 "" diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index b8ea11f3d2f..a759126fcd4 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -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 diff --git a/tests/integration/caching/test_response_cache_model_group_key.py b/tests/integration/caching/test_response_cache_model_group_key.py new file mode 100644 index 00000000000..59565b75d9a --- /dev/null +++ b/tests/integration/caching/test_response_cache_model_group_key.py @@ -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 diff --git a/tests/unit/caching/test_unit_test_caching.py b/tests/unit/caching/test_unit_test_caching.py index 83132b5aa1d..f8a367798db 100644 --- a/tests/unit/caching/test_unit_test_caching.py +++ b/tests/unit/caching/test_unit_test_caching.py @@ -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"]