diff --git a/src/plyngent/agent/budget.py b/src/plyngent/agent/budget.py index cbce059..3693ec0 100644 --- a/src/plyngent/agent/budget.py +++ b/src/plyngent/agent/budget.py @@ -31,27 +31,23 @@ DEFAULT_CONTEXT_MAX_CHARS = DEFAULT_CONTEXT_MAX_TOKENS * 4 def truncate_tool_result( text: str, max_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS, - *, - hint: bool = True, ) -> str: """Cap tool output so huge dumps do not flood model context. - This is the generic agent-loop cap for arbitrary tool results — there is no - resumable source, so no ``TRUNCATE_TOKEN`` is emitted here. Token-capable - sources (``read_file`` / ``fetch`` / ``get_truncated``) embed their own - ``[TRUNCATE_TOKEN: ...]`` marker via :func:`truncate_with_token`, and only - there does the result point the model at ``get_truncated`` (hint and token - always appear together). ``hint=False`` keeps the terse marker for internal - shrinks whose cap is not a user-facing knob (e.g. request-time compacting). + This is the generic agent-loop cap for arbitrary tool results. It embeds a + memory truncate token via :func:`truncate_generic` so the model can always + resume the remainder with ``get_truncated`` — resumable or not (remainders + live in the in-memory store for the process lifetime). The compact marker + ``[Truncated (N chars max.; M omitted). truncate_token=...]`` is emitted + only when output is actually cut. """ if max_chars < 1: return text if len(text) <= max_chars: return text - omitted = len(text) - max_chars - if hint: - return f"{text[:max_chars]}\n...[Truncated ({max_chars} chars max.; {omitted} chars omitted)]" - return f"{text[:max_chars]}\n...[truncated {omitted} characters]" + from plyngent.tools.truncate_token import truncate_generic + + return truncate_generic(text, max_chars) def estimate_message_chars(message: AnyChatMessage) -> int: @@ -119,7 +115,7 @@ def _shrink_tool(message: ToolChatMessage, max_chars: int) -> ToolChatMessage: return message return msgspec.structs.replace( message, - content=truncate_tool_result(message.content, max_chars, hint=False), + content=truncate_tool_result(message.content, max_chars), ) diff --git a/tests/test_agent/test_budget.py b/tests/test_agent/test_budget.py index 9895e5e..a6024ce 100644 --- a/tests/test_agent/test_budget.py +++ b/tests/test_agent/test_budget.py @@ -11,21 +11,17 @@ def test_truncate_tool_result_short() -> None: def test_truncate_tool_result_long() -> None: text = "a" * 50 out = truncate_tool_result(text, 20) - assert out.startswith("a" * 20) assert "Truncated" in out assert "20" in out assert "30" in out - # Generic tool results carry no resumable source → no token, no get_truncated. - assert "get_truncated" not in out - assert "TRUNCATE_TOKEN" not in out + # Generic tool results now embed a memory token too (resumable via get_truncated). + assert "truncate_token=" in out + assert out.startswith(("a" * 10, "a" * 20)) # marker is budgeted -def test_truncate_tool_result_no_hint() -> None: - text = "a" * 50 - out = truncate_tool_result(text, 20, hint=False) - assert out.startswith("a" * 20) - assert "truncated" in out - assert "get_truncated" not in out +def test_truncate_tool_result_fits_no_marker() -> None: + assert truncate_tool_result("tiny", 100) == "tiny" + assert "Truncated" not in truncate_tool_result("tiny", 100) def test_estimate_message_chars() -> None: diff --git a/tests/test_agent/test_compact.py b/tests/test_agent/test_compact.py index e0c73f7..54e9ab7 100644 --- a/tests/test_agent/test_compact.py +++ b/tests/test_agent/test_compact.py @@ -40,21 +40,18 @@ def test_truncate_tool_result_short() -> None: def test_truncate_tool_result_long() -> None: text = "a" * 50 out = truncate_tool_result(text, 20) - assert out.startswith("a" * 20) assert "Truncated" in out assert "20" in out assert "30" in out - # Generic tool results carry no resumable source → no token, no get_truncated. - assert "get_truncated" not in out - assert "TRUNCATE_TOKEN" not in out + # Generic tool results now embed a memory token too (resumable via get_truncated). + assert "truncate_token=" in out def test_truncate_tool_result_no_hint() -> None: text = "a" * 50 - out = truncate_tool_result(text, 20, hint=False) - assert out.startswith("a" * 20) - assert "truncated" in out - assert "get_truncated" not in out + out = truncate_tool_result(text, 20) + assert "Truncated" in out + assert "truncate_token=" in out def test_compact_shrinks_old_tool_results() -> None: @@ -94,7 +91,9 @@ def test_compact_shrinks_old_tool_results() -> None: ) assert isinstance(compacted[2], ToolChatMessage) assert len(compacted[2].content) < original_len - assert "truncated" in compacted[2].content + assert "Truncated" in compacted[2].content + # Compact shrinks carry a token too (resumable via get_truncated). + assert "truncate_token=" in compacted[2].content # Full history unchanged assert isinstance(messages[2], ToolChatMessage) assert len(messages[2].content) == original_len @@ -134,7 +133,7 @@ def test_compact_uses_api_calibration() -> None: # Raw est under budget → no compact no_api = compact_messages_for_request(messages, max_tokens=est + 100) assert isinstance(no_api[1], ToolChatMessage) - assert "truncated" not in no_api[1].content + assert "Truncated" not in no_api[1].content # Calibrate so measured size is 5x → over a mid budget → shrink old tool compacted = compact_messages_for_request( messages, @@ -147,7 +146,7 @@ def test_compact_uses_api_calibration() -> None: assert isinstance(compacted[1], ToolChatMessage) old = messages[1] assert isinstance(old, ToolChatMessage) - assert "truncated" in compacted[1].content or len(compacted[1].content) < len(old.content) + assert "Truncated" in compacted[1].content or len(compacted[1].content) < len(old.content) def _response(message: AssistantChatMessage) -> ChatCompletionResponse: