mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
fix(context_management): count only blocks the model receives and keep list specs unsent
clear_tool_uses now leaves blocks token_counter cannot read, like server tool results and compaction, out of its count instead of counting them as JSON text. No adapter sends them to the model, so an input_tokens trigger no longer fires on tokens the model never sees. On the Responses API path, an OpenAI-style list context_management stays unsent, as it was before, instead of reaching OpenAI's native compaction, where a threshold under 1000 is a 400. On the chat completions path, a clear edit listed before compact_20260112 now runs on the turns after the latest compaction block, as Anthropic does, instead of counting and clearing turns the compaction block drops. The crash test that replaced apply_context_management with a mock is removed.
This commit is contained in:
parent
bcafb0a615
commit
f9d11ccbce
8 changed files with 53 additions and 47 deletions
|
|
@ -160,6 +160,7 @@ async def prepare_context_managed_request(
|
|||
"""Apply client compaction history, then optional context_management polyfill."""
|
||||
from litellm.llms.anthropic.messages.context_management.editors.compact import (
|
||||
apply_client_compaction_block_history,
|
||||
slice_around_compaction_block,
|
||||
)
|
||||
|
||||
# Skip the client-history pre-processing when a ``compact_20260112``
|
||||
|
|
@ -175,7 +176,7 @@ async def prepare_context_managed_request(
|
|||
|
||||
if polyfill_will_run:
|
||||
history_result: PolyfillResult | None = None
|
||||
working_messages: _AnthropicMessages = messages
|
||||
working_messages: _AnthropicMessages = list(slice_around_compaction_block(messages)[0])
|
||||
working_system: _AnthropicSystem = system
|
||||
else:
|
||||
history_result = apply_client_compaction_block_history(
|
||||
|
|
|
|||
|
|
@ -10,7 +10,6 @@ the way the API's 400 does, so the dispatcher can reject the edit, or drop only
|
|||
runs.
|
||||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence, Set
|
||||
from dataclasses import dataclass
|
||||
from functools import reduce
|
||||
|
|
@ -214,17 +213,11 @@ def _select_in_batches(candidates: Sequence[str], tokens: Mapping[str, int], cle
|
|||
return tuple(candidates[:batch_end])
|
||||
|
||||
|
||||
def _countable_block(block: Mapping[str, object]) -> Mapping[str, object]:
|
||||
if block.get("type") in _TOKEN_COUNTER_BLOCK_TYPES:
|
||||
return block
|
||||
return {"type": "text", "text": json.dumps(block, default=str)}
|
||||
|
||||
|
||||
def _countable_message(message: Mapping[str, object]) -> Mapping[str, object]:
|
||||
blocks: Final = _parse(_BLOCKS, message.get("content"))
|
||||
if blocks is None:
|
||||
return message
|
||||
return {**message, "content": [_countable_block(block) for block in blocks]}
|
||||
return {**message, "content": [block for block in blocks if block.get("type") in _TOKEN_COUNTER_BLOCK_TYPES]}
|
||||
|
||||
|
||||
def _count_tokens(
|
||||
|
|
|
|||
|
|
@ -137,7 +137,7 @@ def _find_latest_compaction_index(
|
|||
return None, None
|
||||
|
||||
|
||||
def _slice_around_compaction_block(
|
||||
def slice_around_compaction_block(
|
||||
messages: Sequence[_MsgT],
|
||||
) -> tuple[Sequence[_MsgT | dict[str, object]], dict[str, object] | None]:
|
||||
"""Apply Anthropic's "drop everything before the compaction block" rule.
|
||||
|
|
@ -598,7 +598,7 @@ def apply_client_compaction_block_history(
|
|||
tail is forwarded unchanged (with compaction blocks stripped) so recent
|
||||
turns the summary does not cover are preserved.
|
||||
"""
|
||||
effective_messages, prior_compaction_block = _slice_around_compaction_block(messages)
|
||||
effective_messages, prior_compaction_block = slice_around_compaction_block(messages)
|
||||
if prior_compaction_block is None:
|
||||
return None
|
||||
|
||||
|
|
@ -669,7 +669,7 @@ async def apply_compact_20260112(
|
|||
# opt-in gate below so that even when summarization is disabled we still
|
||||
# strip Anthropic-only ``compaction`` blocks from messages going to
|
||||
# non-Anthropic backends (which would reject them).
|
||||
effective_messages, prior_compaction_block = _slice_around_compaction_block(messages)
|
||||
effective_messages, prior_compaction_block = slice_around_compaction_block(messages)
|
||||
prior_summary_text: Final = prior_compaction_block.get("content") if prior_compaction_block else None
|
||||
augmented_system: str | list[dict[str, object]] | None = system
|
||||
if isinstance(prior_summary_text, str) and prior_summary_text:
|
||||
|
|
|
|||
|
|
@ -50,20 +50,24 @@ AnthropicRequestMessages: TypeAlias = list[AllAnthropicMessageValues] | list[dic
|
|||
BlockRewrite: TypeAlias = Callable[[Mapping[str, object]], Mapping[str, object]]
|
||||
_BLOCK_LIST: Final = TypeAdapter(list[dict[str, object]])
|
||||
_SplitEdits: TypeAlias = tuple[dict[str, object] | None, Sequence[Mapping[str, object]]]
|
||||
AnthropicContextManagementSpec: TypeAlias = dict[str, object] | list[dict[str, object]] | None
|
||||
|
||||
_ADAPTER: Final = LiteLLMAnthropicToResponsesAPIAdapter()
|
||||
|
||||
|
||||
def _split_context_management(
|
||||
spec: dict[str, object] | None, # mutable-ok: the request's spec, handed to normalize_spec
|
||||
spec: AnthropicContextManagementSpec,
|
||||
kwargs: Mapping[str, object],
|
||||
) -> _SplitEdits:
|
||||
"""Return (the edits the Responses API runs natively, the edits the gateway polyfill runs).
|
||||
|
||||
The Responses API compacts natively (``compact_20260112`` becomes its ``compaction``); any other
|
||||
edit, such as ``clear_tool_uses_20250919``, has no Responses API form and runs in the polyfill, as
|
||||
it does on the chat completions path.
|
||||
it does on the chat completions path. An OpenAI-style list spec is not Anthropic's shape and stays
|
||||
unsent, as it was before the polyfill ran here.
|
||||
"""
|
||||
if not isinstance(spec, dict):
|
||||
return None, ()
|
||||
if not spec or context_management_explicitly_dropped(kwargs.get("additional_drop_params")):
|
||||
return spec, ()
|
||||
edits: Final = normalize_spec(spec) or []
|
||||
|
|
@ -332,7 +336,7 @@ class LiteLLMMessagesToResponsesAPIHandler:
|
|||
max_tokens: int,
|
||||
messages: AnthropicRequestMessages,
|
||||
model: str,
|
||||
context_management: dict[str, object] | None = None,
|
||||
context_management: AnthropicContextManagementSpec = None,
|
||||
metadata: dict[str, object] | None = None,
|
||||
output_config: AnthropicOutputConfig | None = None,
|
||||
stop_sequences: list[str] | None = None,
|
||||
|
|
@ -398,7 +402,7 @@ class LiteLLMMessagesToResponsesAPIHandler:
|
|||
max_tokens: int,
|
||||
messages: AnthropicRequestMessages,
|
||||
model: str,
|
||||
context_management: dict[str, object] | None = None,
|
||||
context_management: AnthropicContextManagementSpec = None,
|
||||
metadata: dict[str, object] | None = None,
|
||||
output_config: AnthropicOutputConfig | None = None,
|
||||
stop_sequences: list[str] | None = None,
|
||||
|
|
|
|||
|
|
@ -231,15 +231,10 @@ SERVER_TOOL_TURN: Final = {
|
|||
COMPACTION_TURN: Final = {"role": "assistant", "content": [{"type": "compaction", "content": "summary " * 300}]}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"turn", [pytest.param(SERVER_TOOL_TURN, id="server_tool"), pytest.param(COMPACTION_TURN, id="compaction")]
|
||||
)
|
||||
def test_an_input_tokens_trigger_counts_blocks_the_token_counter_cannot_read(turn: dict[str, object]) -> None:
|
||||
history: Final = _make_history(n_pairs=3)
|
||||
threshold: Final = litellm.token_counter(model=MODEL, messages=history) + 100
|
||||
messages: Final = [history[0], turn, *history[1:]]
|
||||
|
||||
new_messages, applied = apply_clear_tool_uses_20250919(
|
||||
def _clear_above(
|
||||
messages: list[dict[str, object]], threshold: int
|
||||
) -> tuple[list[dict[str, object]], AppliedEdit | None]:
|
||||
return apply_clear_tool_uses_20250919(
|
||||
model=MODEL,
|
||||
messages=messages,
|
||||
tools=None,
|
||||
|
|
@ -251,11 +246,24 @@ def test_an_input_tokens_trigger_counts_blocks_the_token_counter_cannot_read(tur
|
|||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"turn", [pytest.param(SERVER_TOOL_TURN, id="server_tool"), pytest.param(COMPACTION_TURN, id="compaction")]
|
||||
)
|
||||
def test_an_input_tokens_trigger_skips_blocks_no_adapter_sends(turn: dict[str, object]) -> None:
|
||||
history: Final = _make_history(n_pairs=3)
|
||||
sent_tokens: Final = litellm.token_counter(model=MODEL, messages=history)
|
||||
messages: Final = [history[0], turn, *history[1:]]
|
||||
|
||||
untouched, not_applied = _clear_above(messages, sent_tokens + 100)
|
||||
cleared, applied = _clear_above(messages, sent_tokens - 100)
|
||||
|
||||
assert not_applied is None
|
||||
assert untouched == messages
|
||||
assert applied is not None
|
||||
assert applied["cleared_tool_uses"] == 2
|
||||
assert applied["cleared_input_tokens"] > 0
|
||||
assert _cleared_ids(new_messages) == {"toolu_00", "toolu_01"}
|
||||
assert new_messages[1] == turn
|
||||
assert _cleared_ids(cleared) == {"toolu_00", "toolu_01"}
|
||||
assert cleared[1] == turn
|
||||
|
||||
|
||||
def _blocks_of(messages: list[dict[str, object]], block_type: str) -> tuple[dict[str, object], ...]:
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ from litellm.llms.anthropic.messages.context_management.editors.compact import (
|
|||
_augment_system_with_summary,
|
||||
_extract_summary_text,
|
||||
_select_last_user_question,
|
||||
_slice_around_compaction_block,
|
||||
slice_around_compaction_block,
|
||||
_strip_compaction_blocks,
|
||||
apply_client_compaction_block_history,
|
||||
apply_compact_20260112,
|
||||
|
|
@ -153,7 +153,7 @@ def test_applied_edits_for_response_includes_compact_when_block_present():
|
|||
|
||||
def test_slice_around_compaction_block_found():
|
||||
messages = _messages_with_compaction("my summary")
|
||||
sliced, block = _slice_around_compaction_block(messages)
|
||||
sliced, block = slice_around_compaction_block(messages)
|
||||
assert block is not None
|
||||
assert block["type"] == "compaction"
|
||||
assert block["content"] == "my summary"
|
||||
|
|
@ -164,7 +164,7 @@ def test_slice_around_compaction_block_found():
|
|||
|
||||
def test_slice_around_compaction_block_not_found():
|
||||
messages = _simple_messages()
|
||||
sliced, block = _slice_around_compaction_block(messages)
|
||||
sliced, block = slice_around_compaction_block(messages)
|
||||
assert block is None
|
||||
assert sliced is messages # same object, no copy
|
||||
|
||||
|
|
|
|||
|
|
@ -946,27 +946,27 @@ def test_the_sync_path_clears_old_results_too(respx_mock: respx.MockRouter, monk
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_polyfill_crash_is_a_500_before_the_responses_api_call(
|
||||
async def test_an_openai_style_context_management_list_is_still_not_sent_to_the_responses_api(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError
|
||||
from litellm.llms.anthropic.messages.responses_adapters import handler
|
||||
|
||||
upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").respond(
|
||||
status_code=200, json=TEXT_RESPONSE
|
||||
)
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-clear-test")
|
||||
monkeypatch.setattr(handler, "apply_context_management", AsyncMock(side_effect=RuntimeError("editor bug")))
|
||||
with pytest.raises(AnthropicContextManagementError) as raised:
|
||||
await LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler(
|
||||
max_tokens=1024,
|
||||
messages=_three_reads(),
|
||||
model="openai/gpt-5.6-luna",
|
||||
custom_llm_provider="openai",
|
||||
context_management={"edits": [dict(CLEAR_EDIT)]},
|
||||
)
|
||||
assert raised.value.status_code == 500
|
||||
assert not upstream.called
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
served: Final = await LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler(
|
||||
max_tokens=1024,
|
||||
messages=_three_reads(),
|
||||
model="openai/gpt-5.6-luna",
|
||||
custom_llm_provider="openai",
|
||||
context_management=[{"type": "compaction", "compact_threshold": 500}],
|
||||
)
|
||||
|
||||
assert "context_management" not in json.loads(upstream.calls.last.request.content)
|
||||
assert isinstance(served, dict)
|
||||
assert "context_management" not in served
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -649,7 +649,7 @@ async def test_anthropic_view_counts_the_unedited_body_without_an_error_log_when
|
|||
assert not any(record.levelno >= logging.ERROR for record in caplog.records)
|
||||
|
||||
|
||||
async def test_anthropic_view_counts_the_turns_after_a_compaction_block_when_an_edit_before_compact_fails() -> None:
|
||||
async def test_anthropic_view_clears_only_the_turns_after_a_compaction_block_when_clear_runs_before_compact() -> None:
|
||||
clear_before_compact: Final[JsonValue] = {
|
||||
"edits": [
|
||||
{
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue