mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(context_management): count the system prompt toward the clear_tool_uses input_tokens trigger
Anthropic's API compares an input_tokens trigger with the whole prompt, system included; the polyfill counted messages and tools only, so it cleared late by the system prompt's size. compact_20260112 already counts it: share its system_to_text.
This commit is contained in:
parent
b6f91a221b
commit
c6c4541700
3 changed files with 43 additions and 6 deletions
|
|
@ -31,6 +31,7 @@ from ..constants import (
|
|||
)
|
||||
from ..errors import AnthropicContextManagementError
|
||||
from ..placeholders import build_cleared_tool_result_content
|
||||
from .compact import system_to_text
|
||||
|
||||
T = TypeVar("T")
|
||||
# api.anthropic.com, measured 2026-10-19: text results of 1 to 46 bytes are never cleared or counted toward keep
|
||||
|
|
@ -195,11 +196,20 @@ def _select_in_batches(candidates: Sequence[str], tokens: Mapping[str, int], cle
|
|||
return tuple(candidates[:batch_end])
|
||||
|
||||
|
||||
def _with_system(
|
||||
messages: list[dict[str, object]], system: str | list[dict[str, object]] | None
|
||||
) -> list[dict[str, object]]:
|
||||
"""The messages to count: like the API, an ``input_tokens`` trigger counts the system prompt too."""
|
||||
system_text: Final = system_to_text(system)
|
||||
return [{"role": "system", "content": system_text}, *messages] if system_text else messages
|
||||
|
||||
|
||||
def _trigger_met(
|
||||
trigger: dict[str, object],
|
||||
model: str,
|
||||
messages: list[dict[str, object]],
|
||||
tools: list[dict[str, object]] | None,
|
||||
system: str | list[dict[str, object]] | None,
|
||||
counted_tool_uses: int,
|
||||
) -> tuple[bool, int | None]:
|
||||
"""Return (trigger_met, input_tokens if counted for reuse)."""
|
||||
|
|
@ -215,7 +225,7 @@ def _trigger_met(
|
|||
threshold = DEFAULT_INPUT_TOKENS_TRIGGER
|
||||
current_tokens: Final = litellm.token_counter(
|
||||
model=model,
|
||||
messages=messages,
|
||||
messages=_with_system(messages, system),
|
||||
tools=cast(Any, tools),
|
||||
)
|
||||
verbose_logger.debug("context_management polyfill: current_tokens: %s", current_tokens)
|
||||
|
|
@ -332,7 +342,7 @@ def apply_clear_tool_uses_20250919(
|
|||
for tool_use_id, name in _tool_uses_in_order(blocks)
|
||||
if name not in exclude_tools and (tool_use_id not in tiny or _clears_input(name, input_tools))
|
||||
)
|
||||
met, counted_tokens = _trigger_met(trigger, model, messages, tools, len(clearable))
|
||||
met, counted_tokens = _trigger_met(trigger, model, messages, tools, system, len(clearable))
|
||||
if not met:
|
||||
return messages, None
|
||||
|
||||
|
|
@ -366,9 +376,11 @@ def apply_clear_tool_uses_20250919(
|
|||
tokens_before: Final = (
|
||||
counted_tokens
|
||||
if counted_tokens is not None
|
||||
else litellm.token_counter(model=model, messages=messages, tools=cast(Any, tools))
|
||||
else litellm.token_counter(model=model, messages=_with_system(messages, system), tools=cast(Any, tools))
|
||||
)
|
||||
tokens_after: Final = litellm.token_counter(
|
||||
model=model, messages=_with_system(final_messages, system), tools=cast(Any, tools)
|
||||
)
|
||||
tokens_after: Final = litellm.token_counter(model=model, messages=final_messages, tools=cast(Any, tools))
|
||||
cleared_input_tokens: Final = max(tokens_before - tokens_after, 0)
|
||||
|
||||
return final_messages, {
|
||||
|
|
|
|||
|
|
@ -320,13 +320,13 @@ def _count_effective_tokens(
|
|||
content: Final = compaction_block.get("content") or ""
|
||||
if content:
|
||||
total += litellm.token_counter(model=model, text=content)
|
||||
system_text: Final = _system_to_text(system)
|
||||
system_text: Final = system_to_text(system)
|
||||
if system_text:
|
||||
total += litellm.token_counter(model=model, text=system_text)
|
||||
return total
|
||||
|
||||
|
||||
def _system_to_text(
|
||||
def system_to_text(
|
||||
system: str | list[dict[str, object]] | None,
|
||||
) -> str:
|
||||
"""Flatten an Anthropic-style ``system`` value into a single string for
|
||||
|
|
|
|||
|
|
@ -690,6 +690,31 @@ def test_tool_uses_trigger_does_not_count_excluded_tools(trigger_value: int, tri
|
|||
assert (applied is not None) == triggers
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"system",
|
||||
["Reference notes. " * 400, [{"type": "text", "text": "Reference notes. " * 400}]],
|
||||
)
|
||||
def test_input_tokens_trigger_counts_the_system_prompt(system: str | list[dict[str, object]]) -> None:
|
||||
"""Like the API, the trigger compares the whole prompt: a system prompt can push it over."""
|
||||
messages = _make_history(n_pairs=3)
|
||||
history_tokens: Final = litellm.token_counter(model=MODEL, messages=messages)
|
||||
edit_spec: Final = {
|
||||
"type": "clear_tool_uses_20250919",
|
||||
"trigger": {"type": "input_tokens", "value": history_tokens + 100},
|
||||
"keep": {"type": "tool_uses", "value": 1},
|
||||
}
|
||||
_, without_system = apply_clear_tool_uses_20250919(
|
||||
model=MODEL, messages=deepcopy(messages), tools=None, system=None, edit_spec=edit_spec
|
||||
)
|
||||
_, with_system = apply_clear_tool_uses_20250919(
|
||||
model=MODEL, messages=deepcopy(messages), tools=None, system=system, edit_spec=edit_spec
|
||||
)
|
||||
assert without_system is None
|
||||
assert with_system is not None
|
||||
assert with_system["cleared_tool_uses"] == 2
|
||||
assert with_system["cleared_input_tokens"] < history_tokens
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("clear_tool_inputs", "triggers"), [(False, False), (True, True)])
|
||||
def test_tool_uses_trigger_counts_tiny_results_only_when_their_inputs_are_cleared(
|
||||
clear_tool_inputs: bool, triggers: bool
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue