diff --git a/python/packages/lab/tau2/agent_framework_lab_tau2/_sliding_window.py b/python/packages/lab/tau2/agent_framework_lab_tau2/_sliding_window.py index 3a6157f3ef0..1218f047971 100644 --- a/python/packages/lab/tau2/agent_framework_lab_tau2/_sliding_window.py +++ b/python/packages/lab/tau2/agent_framework_lab_tau2/_sliding_window.py @@ -100,16 +100,16 @@ def _get_token_count(self, messages: list[Message]) -> int: # Simple text message total_tokens += self._estimate_any_object_token_count(msg.text) - if total_tokens > self.max_tokens / 2: + if total_tokens > self.max_tokens: + logger.opt(colors=True).warning( + f"Total tokens {total_tokens} is over max tokens {self.max_tokens}. Will truncate messages." + ) + elif total_tokens > self.max_tokens / 2: logger.opt(colors=True).warning( f"Total tokens {total_tokens} is " f"{total_tokens / self.max_tokens * 100:.0f}% " f"of max tokens {self.max_tokens}" ) - elif total_tokens > self.max_tokens: - logger.opt(colors=True).warning( - f"Total tokens {total_tokens} is over max tokens {self.max_tokens}. Will truncate messages." - ) return total_tokens diff --git a/python/packages/lab/tau2/tests/test_sliding_window.py b/python/packages/lab/tau2/tests/test_sliding_window.py index 8da46e86958..1a45485146f 100644 --- a/python/packages/lab/tau2/tests/test_sliding_window.py +++ b/python/packages/lab/tau2/tests/test_sliding_window.py @@ -211,3 +211,39 @@ async def test_real_world_scenario(): token_count = provider._get_token_count(truncated) assert token_count <= provider.max_tokens * 1.1 + + +def test_get_token_count_warns_over_max(): + """The over-max warning must win over the half-of-max warning (#9338).""" + import io + + from loguru import logger + + provider = SlidingWindowHistoryProvider(max_tokens=10) + messages = [ + Message( + role="user", + contents=[ + Content.from_text( + text=( + "hello world, this message is intentionally long " + "so it clearly exceeds the tiny ten token budget" + ) + ) + ], + ) + ] + + total = provider._get_token_count(messages) + assert total > provider.max_tokens + + buffer = io.StringIO() + handler_id = logger.add(buffer, format="{message}") + try: + provider._get_token_count(messages) + finally: + logger.remove(handler_id) + + logged = buffer.getvalue() + assert "over max tokens" in logged + assert "Will truncate messages" in logged