mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
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.
This commit is contained in:
parent
3ecc9efb7e
commit
de7ad25216
4 changed files with 52 additions and 30 deletions
|
|
@ -461,16 +461,14 @@ class Cache:
|
|||
litellm_params.get("litellm_metadata") or {},
|
||||
litellm_params.get("metadata") or {},
|
||||
)
|
||||
model_group: Final[str | None] = next(
|
||||
(source["model_group"] for source in metadata_sources if source.get("model_group")), None
|
||||
)
|
||||
caching_group: Final = next(
|
||||
(group for source in metadata_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}"
|
||||
model_group_source: Final = 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
|
||||
|
||||
def _get_caching_group(self, metadata: dict, model_group: str | None) -> str | None:
|
||||
caching_groups: Final[list | None] = metadata.get("caching_groups", [])
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -38,32 +38,32 @@ from tests.integration.caching.response_cache_case import (
|
|||
scripted,
|
||||
)
|
||||
|
||||
Ask: TypeAlias = Callable[[Gateway, str, str, Mapping[str, str]], str]
|
||||
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, str]) -> dict[str, JsonValue] | None:
|
||||
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, str]) -> str:
|
||||
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, str]) -> str:
|
||||
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, str]) -> str:
|
||||
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,
|
||||
|
|
@ -88,15 +88,15 @@ class Group:
|
|||
|
||||
|
||||
ENDPOINTS: Final = {
|
||||
"chat_completions": Endpoint(CHAT_MODEL, "openai/gpt-4o", CHAT, _chat),
|
||||
"responses": Endpoint(CHAT_MODEL, "openai/gpt-4o", RESPONSE, _responses),
|
||||
"messages": Endpoint(MESSAGES_MODEL, "anthropic/claude-sonnet-4-5", MESSAGE, _messages),
|
||||
"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) -> Group:
|
||||
def _group(scenario: Scenario, endpoint: Endpoint, name: str | None = None) -> Group:
|
||||
handle: Final = scripted(scenario, json_response(endpoint.answer))
|
||||
return Group(scenario.model(model=endpoint.provider_model, api_base=handle.api_base()), handle)
|
||||
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())
|
||||
|
|
@ -118,17 +118,24 @@ def test_identical_requests_to_two_groups_on_one_model_each_reach_their_own_depl
|
|||
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
|
||||
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:
|
||||
target: Final = _group(scenario, endpoint)
|
||||
caller: Final = _group(scenario, endpoint)
|
||||
planted: Final = endpoint.ask(gateway, caller.name, text, {"model_group": target.name})
|
||||
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, {})
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ load_dotenv()
|
|||
import asyncio
|
||||
import hashlib
|
||||
import random
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -279,8 +280,8 @@ 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 = Cache()
|
||||
request = {"model": "openai/gpt-5.6", "input": "search this text"}
|
||||
cache: Final = Cache()
|
||||
request: Final = {"model": "openai/gpt-5.6", "input": "search this text"}
|
||||
|
||||
assert (
|
||||
cache._get_model_param_value(
|
||||
|
|
@ -307,6 +308,20 @@ def test_get_model_param_value_prefers_the_router_group_over_a_caller_model_grou
|
|||
== "group-a"
|
||||
)
|
||||
|
||||
|
||||
@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):
|
||||
cache: Final = Cache()
|
||||
request: Final = {
|
||||
"model": "openai/gpt-5.6",
|
||||
"input": "search this text",
|
||||
"litellm_metadata": {"model_group": "sales", "caching_groups": None},
|
||||
"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