mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-30 06:43:27 +00:00
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:
parent
41602e318d
commit
ebf84f8649
4 changed files with 326 additions and 14 deletions
|
|
@ -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 ""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
192
tests/integration/caching/test_response_cache_model_group_key.py
Normal file
192
tests/integration/caching/test_response_cache_model_group_key.py
Normal 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
|
||||
|
|
@ -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"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue