mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-01 02:02:20 +00:00
fix(compact_20260112): propagate allowed_model_region to summary subrequest
The router enforces region restrictions by reading allowed_model_region from top-level request kwargs (Router._common_checks_available_deployment), but the compact_20260112 summary subrequest only forwarded litellm_metadata. A region-restricted caller could trigger compaction and have their conversation summarized by a deployment outside the permitted region. Extract allowed_model_region from user_api_key_auth and pass it through _call_summary_model as a top-level kwarg so the router applies the same region constraints the parent request would.
This commit is contained in:
parent
0adbb24633
commit
d41ab61409
2 changed files with 94 additions and 0 deletions
|
|
@ -525,6 +525,7 @@ async def _call_summary_model(
|
|||
summary_messages: List[Dict[str, Any]],
|
||||
metadata: Dict[str, Any],
|
||||
llm_router: Any,
|
||||
allowed_model_region: Optional[str] = None,
|
||||
) -> Any:
|
||||
"""Invoke the configured summary model.
|
||||
|
||||
|
|
@ -541,12 +542,19 @@ async def _call_summary_model(
|
|||
# post-call spend hooks read for budget attribution. The provider-level
|
||||
# ``metadata`` kwarg corresponds to the upstream API request body and would
|
||||
# not flow into spend tracking.
|
||||
# ``allowed_model_region`` must travel as a top-level kwarg because the
|
||||
# router enforces region restrictions by reading ``request_kwargs`` directly
|
||||
# (see ``Router._common_checks_available_deployment``); without this the
|
||||
# summary subrequest could be routed to a deployment outside the caller's
|
||||
# permitted region.
|
||||
call_kwargs: Dict[str, Any] = {
|
||||
"model": summary_model,
|
||||
"messages": summary_messages,
|
||||
"max_tokens": COMPACT_SUMMARY_MAX_TOKENS,
|
||||
"litellm_metadata": metadata,
|
||||
}
|
||||
if allowed_model_region is not None:
|
||||
call_kwargs["allowed_model_region"] = allowed_model_region
|
||||
if llm_router is not None and hasattr(llm_router, "acompletion"):
|
||||
return await llm_router.acompletion(**call_kwargs)
|
||||
return await litellm.acompletion(**call_kwargs)
|
||||
|
|
@ -754,6 +762,7 @@ async def apply_compact_20260112( # noqa: PLR0915
|
|||
effective_messages, prompt, system=augmented_system
|
||||
)
|
||||
propagated_metadata = _propagate_metadata(litellm_metadata)
|
||||
allowed_model_region = getattr(user_api_key_auth, "allowed_model_region", None)
|
||||
|
||||
try:
|
||||
response = await _call_summary_model(
|
||||
|
|
@ -761,6 +770,7 @@ async def apply_compact_20260112( # noqa: PLR0915
|
|||
summary_messages=summary_messages,
|
||||
metadata=propagated_metadata,
|
||||
llm_router=llm_router,
|
||||
allowed_model_region=allowed_model_region,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.warning("compact_20260112: summary call failed: %s", e)
|
||||
|
|
|
|||
|
|
@ -1194,6 +1194,90 @@ async def test_summary_model_allowed_when_no_user_api_key_auth():
|
|||
assert result.compaction_block is not None
|
||||
|
||||
|
||||
async def test_summary_call_propagates_allowed_model_region():
|
||||
"""``allowed_model_region`` from ``user_api_key_auth`` is propagated to the
|
||||
summary subrequest as a top-level kwarg so the router applies the same
|
||||
region restriction the parent request would.
|
||||
"""
|
||||
messages = _simple_messages()
|
||||
mock_call = AsyncMock(return_value=_make_mock_response("<summary>ok</summary>"))
|
||||
|
||||
auth = _fake_user_api_key_auth(key_models=["all-proxy-models"])
|
||||
auth.allowed_model_region = "eu"
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting",
|
||||
return_value="claude-haiku-4-5",
|
||||
),
|
||||
patch("litellm.token_counter", return_value=200_000),
|
||||
patch(
|
||||
"litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model",
|
||||
mock_call,
|
||||
),
|
||||
):
|
||||
await apply_compact_20260112(
|
||||
model=MODEL,
|
||||
messages=messages,
|
||||
tools=None,
|
||||
system=None,
|
||||
edit_spec=_EDIT_SPEC_DEFAULT,
|
||||
user_api_key_auth=auth,
|
||||
)
|
||||
|
||||
mock_call.assert_awaited_once()
|
||||
assert mock_call.await_args.kwargs.get("allowed_model_region") == "eu"
|
||||
|
||||
|
||||
async def test_summary_call_omits_allowed_model_region_when_unset():
|
||||
"""Callers without a region restriction must not get an ``allowed_model_region=None``
|
||||
kwarg, which would otherwise force the router to evaluate region filtering.
|
||||
"""
|
||||
from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
|
||||
_call_summary_model,
|
||||
)
|
||||
|
||||
captured_kwargs: dict = {}
|
||||
|
||||
class _FakeRouter:
|
||||
async def acompletion(self, **kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return _make_mock_response("<summary>x</summary>")
|
||||
|
||||
await _call_summary_model(
|
||||
summary_model="claude-haiku-4-5",
|
||||
summary_messages=[{"role": "user", "content": "hi"}],
|
||||
metadata={},
|
||||
llm_router=_FakeRouter(),
|
||||
)
|
||||
|
||||
assert "allowed_model_region" not in captured_kwargs
|
||||
|
||||
|
||||
async def test_summary_call_forwards_allowed_model_region_when_set():
|
||||
"""When the caller is region-restricted, the kwarg reaches the router."""
|
||||
from litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact import (
|
||||
_call_summary_model,
|
||||
)
|
||||
|
||||
captured_kwargs: dict = {}
|
||||
|
||||
class _FakeRouter:
|
||||
async def acompletion(self, **kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return _make_mock_response("<summary>x</summary>")
|
||||
|
||||
await _call_summary_model(
|
||||
summary_model="claude-haiku-4-5",
|
||||
summary_messages=[{"role": "user", "content": "hi"}],
|
||||
metadata={},
|
||||
llm_router=_FakeRouter(),
|
||||
allowed_model_region="eu",
|
||||
)
|
||||
|
||||
assert captured_kwargs.get("allowed_model_region") == "eu"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dispatcher integration: compact_20260112 via apply_context_management
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue