From 9821aab0d5265a4e6fabc86c1b04ecf4b0eedcb2 Mon Sep 17 00:00:00 2001 From: David SF <64162682+dsfaccini@users.noreply.github.com> Date: Mon, 15 Jun 2026 15:34:19 -0500 Subject: [PATCH] feat(compaction): add ClampOversizedMessages for runaway generations (#286) A single runaway model response (degenerate whitespace) or giant tool-call payload can produce one part large enough to blow the next request past the provider context cap. No existing strategy can reach it: SlidingWindow drops the oldest messages, ClearToolResults only touches tool results, and feeding the history to SummarizingCompaction hits the same cap. ClampOversizedMessages truncates the offending ModelResponse part in place, keeping a head/tail slice with a marker. Covers response text (critical case) and tool-call args (clamp_tool_call_args, default on); request-side parts stay out of scope. Composes as the first zero-LLM tier of TieredCompaction. Closes #285 Co-authored-by: David SF Co-authored-by: Claude Opus 4.8 (1M context) --- .../experimental/compaction/README.md | 59 ++++++ .../experimental/compaction/__init__.py | 2 + .../compaction/_clamp_oversized_messages.py | 174 ++++++++++++++++++ .../experimental/compaction/_shared.py | 10 + .../compaction/test_compaction.py | 170 +++++++++++++++++ 5 files changed, 415 insertions(+) create mode 100644 pydantic_ai_harness/experimental/compaction/_clamp_oversized_messages.py diff --git a/pydantic_ai_harness/experimental/compaction/README.md b/pydantic_ai_harness/experimental/compaction/README.md index d2d0eea..485fa0f 100644 --- a/pydantic_ai_harness/experimental/compaction/README.md +++ b/pydantic_ai_harness/experimental/compaction/README.md @@ -31,6 +31,7 @@ provider rejects an orphaned pair. The zero-LLM strategies never call a model. | Capability | Cost | What it does | Reach for it when | |---|---|---|---| +| `ClampOversizedMessages` | zero-LLM | Head/tail-truncates a single oversized part (response text, tool-call args) | One runaway generation blew past the context cap and no other strategy can reach it | | `SlidingWindow` | zero-LLM | Drops the oldest whole messages down to a tail | You only need the recent turns and can discard old context entirely | | `ClearToolResults` | zero-LLM | Blanks the content of old tool *results* in place, keeping the last `keep_pairs` | Tool outputs dominate context and can be re-fetched on demand (the cheap first tier) | | `DeduplicateFileReads` | zero-LLM | Blanks every file read superseded by a newer read of the same file | The agent re-reads files and only the latest version matters | @@ -44,6 +45,64 @@ Every size-based strategy triggers on `max_messages` and/or `max_tokens` (estima use a ~4-chars-per-token heuristic by default; pass a `tokenizer` callable (e.g. `tiktoken`) for accuracy. `DeduplicateFileReads` runs on every request when no trigger is set (it is cheap and near-lossless). `TieredCompaction` triggers and stops on a single `target_tokens` budget. +`ClampOversizedMessages` triggers per *part* (`max_part_tokens` / `max_part_chars`), not on the +whole history -- the failure it targets is one oversized part, not a large total. + +## `ClampOversizedMessages`: surviving a runaway generation + +A single model response of repeated whitespace, or a single tool call with a giant payload, can +produce one part so large the *next* request exceeds the provider's context cap. None of the other +strategies can reach it: `SlidingWindow` drops the oldest messages but the offender is the newest; +`ClearToolResults` only touches tool *results*; `LimitWarner` never edits history; and feeding the +history to `SummarizingCompaction` hits the same cap. + +`ClampOversizedMessages` truncates the offending part in place, keeping a head slice and a tail slice +with a `[clamped: removed N of M characters]` marker between them. Degenerate generations are +low-entropy repetition, so a head/tail slice loses little. + +```python +from pydantic_ai import Agent +from pydantic_ai_harness.experimental.compaction import ClampOversizedMessages + +agent = Agent( + 'openai:gpt-4o', + capabilities=[ClampOversizedMessages(max_part_tokens=50_000, keep_head_chars=2_000, keep_tail_chars=2_000)], +) +``` + +A part is clamped only when it is oversized *and* the clamp actually shrinks it, so keep +`keep_head_chars + keep_tail_chars` well below your per-part threshold. + +It clamps two kinds of part inside each `ModelResponse`: + +- **Response text** (`TextPart`) -- the critical case, a runaway model-response text part. +- **Tool-call args** (`ToolCallPart`), when `clamp_tool_call_args=True` (default) -- the same failure + shape for a giant payload (e.g. a runaway `write_plan`). The args are replaced with a small JSON + object `{"_clamped": "..."}` so they stay valid function arguments; the original call + already executed, so this only shrinks the history copy. Set `clamp_tool_call_args=False` to clamp + response text only. + +Request-side parts (user prompts, tool *returns*, system prompts) are deliberately out of scope: +user input should not be silently rewritten, and oversized tool returns are the job of +`ClearToolResults`. + +Use it as the first tier of `TieredCompaction`, before `ClearToolResults`: + +```python +from pydantic_ai_harness.experimental.compaction import ( + ClampOversizedMessages, + ClearToolResults, + TieredCompaction, +) + +TieredCompaction( + tiers=[ + ClampOversizedMessages(max_part_tokens=50_000), + ClearToolResults(max_tokens=1, keep_pairs=3), + ], + target_tokens=120_000, +) +``` ## Cost: why summarization is the last resort diff --git a/pydantic_ai_harness/experimental/compaction/__init__.py b/pydantic_ai_harness/experimental/compaction/__init__.py index 715a3f7..c501271 100644 --- a/pydantic_ai_harness/experimental/compaction/__init__.py +++ b/pydantic_ai_harness/experimental/compaction/__init__.py @@ -5,6 +5,7 @@ Each capability lives in its own module; shared utilities (token estimation, the """ from pydantic_ai_harness.experimental._warn import warn_experimental +from pydantic_ai_harness.experimental.compaction._clamp_oversized_messages import ClampOversizedMessages from pydantic_ai_harness.experimental.compaction._clear_tool_results import ClearToolResults from pydantic_ai_harness.experimental.compaction._deduplicate_file_reads import DeduplicateFileReads from pydantic_ai_harness.experimental.compaction._limit_warner import LimitWarner, WarningKind @@ -16,6 +17,7 @@ from pydantic_ai_harness.experimental.compaction._tiered_compaction import Tiere warn_experimental('compaction') __all__ = [ + 'ClampOversizedMessages', 'ClearToolResults', 'CompactionStrategy', 'DeduplicateFileReads', diff --git a/pydantic_ai_harness/experimental/compaction/_clamp_oversized_messages.py b/pydantic_ai_harness/experimental/compaction/_clamp_oversized_messages.py new file mode 100644 index 0000000..cf25f5f --- /dev/null +++ b/pydantic_ai_harness/experimental/compaction/_clamp_oversized_messages.py @@ -0,0 +1,174 @@ +"""`ClampOversizedMessages` -- zero-cost head/tail truncation of a single oversized part.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING + +from pydantic_ai._run_context import AgentDepsT +from pydantic_ai.capabilities import AbstractCapability +from pydantic_ai.messages import ( + ModelMessage, + ModelResponse, + ModelResponsePart, + TextPart, + ToolCallPart, +) +from pydantic_ai.tools import RunContext + +from pydantic_ai_harness.experimental.compaction._shared import estimate_text_tokens + +if TYPE_CHECKING: + from pydantic_ai.models import ModelRequestContext + + +_CLAMP_MARKER = '\n[clamped: removed {removed} of {original} characters]\n' +"""Inserted between the head and tail slices of a clamped part. ``{removed}`` and ``{original}`` +are filled with character counts.""" + +_CLAMP_ARGS_KEY = '_clamped' +"""Key of the single-entry object a clamped `ToolCallPart`'s args are replaced with. + +Args stay a JSON object (not a bare marker string) so `args_as_json_str()` emits valid function +arguments for the provider.""" + + +@dataclass +class ClampOversizedMessages(AbstractCapability[AgentDepsT]): + """Zero-cost head/tail truncation of any single oversized message part. + + A runaway generation -- a model response of repeated whitespace, a giant tool-call + payload -- can produce one part so large the next request exceeds the provider's context + cap. The size-based strategies cannot help: `SlidingWindow` drops the *oldest* messages + (the offender is the newest), `ClearToolResults` only touches tool *results*, and feeding + the history to `SummarizingCompaction` hits the same cap. This strategy truncates the + offending part in place: it keeps a head slice and a tail slice and inserts a marker for + the removed middle. Degenerate generations are low-entropy repetition, so a head/tail + slice loses little. No LLM calls are made. + + What it clamps, in each `ModelResponse`: + + - `TextPart` content (the critical case -- a runaway model-response text part). + - `ToolCallPart` args, when `clamp_tool_call_args` is set (the same failure shape for a + giant tool-call payload). The args are replaced with a small JSON object so they stay + valid function arguments; the original call already executed, so this only shrinks the + history copy. + + Request-side parts (user prompts, tool returns, system prompts) are out of scope: user + input should not be silently rewritten, and oversized tool *returns* are the job of + `ClearToolResults`. + + Clamping rewrites message content, so it invalidates the provider's prompt cache from the + clamped message onward. That is unavoidable here -- the alternative is a failed request. + + A part is clamped only when it is oversized *and* the clamp actually shrinks it, so set + `keep_head_chars` + `keep_tail_chars` well below your per-part threshold. + + Composes as the first tier of a `TieredCompaction` (run it before `ClearToolResults`): + it is the only zero-LLM way to keep a run alive after a runaway generation. + + Example: + ```python + from pydantic_ai import Agent + from pydantic_ai_harness.experimental.compaction import ClampOversizedMessages + + agent = Agent( + 'openai:gpt-4o', + capabilities=[ClampOversizedMessages(max_part_tokens=50_000)], + ) + ``` + """ + + max_part_tokens: int | None = None + """Clamp a part whose estimated token count exceeds this value. ``None`` disables this trigger.""" + + max_part_chars: int | None = None + """Clamp a part whose character count exceeds this value. ``None`` disables this trigger.""" + + keep_head_chars: int = 2_000 + """Characters of the part's head to retain.""" + + keep_tail_chars: int = 2_000 + """Characters of the part's tail to retain.""" + + clamp_tool_call_args: bool = True + """When ``True``, also clamp oversized `ToolCallPart` args, not just response text.""" + + tokenizer: Callable[[str], int] | None = None + """Optional tokenizer for accurate token counting. + + A callable that returns the token count for a given string. + When ``None``, uses a ~4 characters-per-token heuristic. + """ + + def __post_init__(self) -> None: + if self.max_part_tokens is None and self.max_part_chars is None: + raise ValueError('At least one of max_part_tokens or max_part_chars must be set.') + if self.max_part_tokens is not None and self.max_part_tokens < 1: + raise ValueError('max_part_tokens must be positive.') + if self.max_part_chars is not None and self.max_part_chars < 1: + raise ValueError('max_part_chars must be positive.') + if self.keep_head_chars < 0: + raise ValueError('keep_head_chars must be non-negative.') + if self.keep_tail_chars < 0: + raise ValueError('keep_tail_chars must be non-negative.') + + def _is_oversized(self, text: str) -> bool: + if self.max_part_chars is not None and len(text) > self.max_part_chars: + return True + if self.max_part_tokens is not None and estimate_text_tokens(text, self.tokenizer) > self.max_part_tokens: + return True + return False + + def _clamp(self, text: str) -> str | None: + """Return the head/tail-clamped form of *text*, or ``None`` if it would not shrink.""" + if not self._is_oversized(text): + return None + head = text[: self.keep_head_chars] + tail = text[len(text) - self.keep_tail_chars :] if self.keep_tail_chars else '' + removed = len(text) - len(head) - len(tail) + clamped = head + _CLAMP_MARKER.format(removed=removed, original=len(text)) + tail + if len(clamped) >= len(text): + return None + return clamped + + async def compact( + self, + messages: list[ModelMessage], + ctx: RunContext[AgentDepsT], + ) -> list[ModelMessage]: + """Clamp every oversized response text part (and tool-call args, if enabled).""" + out: list[ModelMessage] = [] + for msg in messages: + if not isinstance(msg, ModelResponse): + out.append(msg) + continue + + new_parts: list[ModelResponsePart] = [] + changed = False + for part in msg.parts: + if isinstance(part, TextPart): + clamped = self._clamp(part.content) + if clamped is not None: + new_parts.append(replace(part, content=clamped)) + changed = True + continue + elif isinstance(part, ToolCallPart) and self.clamp_tool_call_args: + clamped = self._clamp(part.args_as_json_str()) + if clamped is not None: + new_parts.append(replace(part, args={_CLAMP_ARGS_KEY: clamped})) + changed = True + continue + new_parts.append(part) + out.append(replace(msg, parts=new_parts) if changed else msg) + return out + + async def before_model_request( + self, + ctx: RunContext[AgentDepsT], + request_context: ModelRequestContext, + ) -> ModelRequestContext: + """Clamp any oversized response part before the request is sent.""" + request_context.messages = await self.compact(list(request_context.messages), ctx) + return request_context diff --git a/pydantic_ai_harness/experimental/compaction/_shared.py b/pydantic_ai_harness/experimental/compaction/_shared.py index 4116e5e..d4099b1 100644 --- a/pydantic_ai_harness/experimental/compaction/_shared.py +++ b/pydantic_ai_harness/experimental/compaction/_shared.py @@ -69,6 +69,16 @@ def _user_prompt_text_for_counting(part: UserPromptPart) -> str: return ''.join(texts) +def estimate_text_tokens(text: str, tokenizer: Callable[[str], int] | None = None) -> int: + """Approximate the token count of a single string. + + Uses *tokenizer* when given, otherwise the ~4 characters-per-token heuristic. + """ + if tokenizer is not None: + return tokenizer(text) + return len(text) // _CHARS_PER_TOKEN + + def estimate_token_count( messages: Sequence[ModelMessage], tokenizer: Callable[[str], int] | None = None, diff --git a/tests/experimental/compaction/test_compaction.py b/tests/experimental/compaction/test_compaction.py index ff3bb6a..230784c 100644 --- a/tests/experimental/compaction/test_compaction.py +++ b/tests/experimental/compaction/test_compaction.py @@ -21,6 +21,7 @@ from pydantic_ai.models import ModelRequestContext, ModelRequestParameters from pydantic_ai.usage import RunUsage from pydantic_ai_harness.experimental.compaction import ( + ClampOversizedMessages, ClearToolResults, DeduplicateFileReads, LimitWarner, @@ -29,6 +30,10 @@ from pydantic_ai_harness.experimental.compaction import ( TieredCompaction, estimate_token_count, ) +from pydantic_ai_harness.experimental.compaction._clamp_oversized_messages import ( + _CLAMP_ARGS_KEY, + _CLAMP_MARKER, +) from pydantic_ai_harness.experimental.compaction._shared import ( _is_safe_cutoff, find_first_user_message, @@ -1796,6 +1801,159 @@ class TestSummarizingCompactionModel: assert heading in _DEFAULT_SUMMARY_PROMPT +class TestClampOversizedMessages: + def test_validation_no_trigger(self): + with pytest.raises(ValueError, match='max_part_tokens or max_part_chars'): + ClampOversizedMessages() + + def test_validation_negative_max_part_tokens(self): + with pytest.raises(ValueError, match='max_part_tokens must be positive'): + ClampOversizedMessages(max_part_tokens=0) + + def test_validation_negative_max_part_chars(self): + with pytest.raises(ValueError, match='max_part_chars must be positive'): + ClampOversizedMessages(max_part_chars=0) + + def test_validation_negative_keep_head(self): + with pytest.raises(ValueError, match='keep_head_chars must be non-negative'): + ClampOversizedMessages(max_part_chars=10, keep_head_chars=-1) + + def test_validation_negative_keep_tail(self): + with pytest.raises(ValueError, match='keep_tail_chars must be non-negative'): + ClampOversizedMessages(max_part_chars=10, keep_tail_chars=-1) + + @pytest.mark.anyio + async def test_clamps_oversized_response_text(self): + text = 'H' * 50 + ' ' * 5_000 + 'T' * 50 + cap = ClampOversizedMessages(max_part_chars=1_000, keep_head_chars=50, keep_tail_chars=50) + messages: list[ModelMessage] = [_assistant(text)] + result = await cap.compact(messages, _make_ctx()) + + clamped = result[0] + assert isinstance(clamped, ModelResponse) + part = clamped.parts[0] + assert isinstance(part, TextPart) + assert part.content.startswith('H' * 50) + assert part.content.endswith('T' * 50) + assert '[clamped: removed' in part.content + assert len(part.content) < len(text) + + @pytest.mark.anyio + async def test_token_trigger_uses_heuristic(self): + text = 'x' * 4_000 # ~1000 tokens at the 4-chars heuristic. + cap = ClampOversizedMessages(max_part_tokens=100, keep_head_chars=20, keep_tail_chars=20) + result = await cap.compact([_assistant(text)], _make_ctx()) + part = result[0].parts[0] + assert isinstance(part, TextPart) + assert '[clamped: removed' in part.content + + @pytest.mark.anyio + async def test_token_trigger_uses_tokenizer(self): + text = 'word ' * 1_000 + cap = ClampOversizedMessages( + max_part_tokens=100, + keep_head_chars=20, + keep_tail_chars=20, + tokenizer=lambda s: len(s.split()), + ) + result = await cap.compact([_assistant(text)], _make_ctx()) + part = result[0].parts[0] + assert isinstance(part, TextPart) + assert '[clamped: removed' in part.content + + @pytest.mark.anyio + async def test_small_text_untouched(self): + cap = ClampOversizedMessages(max_part_chars=100_000, max_part_tokens=100_000) + messages: list[ModelMessage] = [_assistant('short')] + result = await cap.compact(messages, _make_ctx()) + # Nothing oversized -> the message object is returned unchanged. + assert result[0] is messages[0] + + @pytest.mark.anyio + async def test_keep_tail_zero(self): + text = 'A' * 5_000 + cap = ClampOversizedMessages(max_part_chars=1_000, keep_head_chars=100, keep_tail_chars=0) + result = await cap.compact([_assistant(text)], _make_ctx()) + part = result[0].parts[0] + assert isinstance(part, TextPart) + assert part.content.startswith('A' * 100) + assert part.content.endswith(']\n') + + @pytest.mark.anyio + async def test_clamp_skipped_when_not_smaller(self): + # Oversized by the token trigger, but keep slices exceed the text length, so + # clamping would not shrink it -- leave it untouched. + text = 'y' * 100 + cap = ClampOversizedMessages(max_part_tokens=1, keep_head_chars=2_000, keep_tail_chars=2_000) + messages: list[ModelMessage] = [_assistant(text)] + result = await cap.compact(messages, _make_ctx()) + assert result[0] is messages[0] + + @pytest.mark.anyio + async def test_clamps_oversized_tool_call_args(self): + big = 'p' * 5_000 + call = ModelResponse(parts=[ToolCallPart(tool_name='write_plan', args=big, tool_call_id='c1')]) + cap = ClampOversizedMessages(max_part_chars=1_000, keep_head_chars=50, keep_tail_chars=50) + result = await cap.compact([call], _make_ctx()) + + part = result[0].parts[0] + assert isinstance(part, ToolCallPart) + assert isinstance(part.args, dict) + assert _CLAMP_ARGS_KEY in part.args + assert '[clamped: removed' in part.args[_CLAMP_ARGS_KEY] + assert part.tool_call_id == 'c1' + + @pytest.mark.anyio + async def test_small_tool_call_args_untouched(self): + call = ModelResponse(parts=[ToolCallPart(tool_name='t', args='{"a": 1}', tool_call_id='c1')]) + cap = ClampOversizedMessages(max_part_chars=1_000) + messages: list[ModelMessage] = [call] + result = await cap.compact(messages, _make_ctx()) + assert result[0] is messages[0] + + @pytest.mark.anyio + async def test_tool_call_args_not_clamped_when_disabled(self): + big = 'p' * 5_000 + call = ModelResponse(parts=[ToolCallPart(tool_name='write_plan', args=big, tool_call_id='c1')]) + cap = ClampOversizedMessages(max_part_chars=1_000, clamp_tool_call_args=False) + messages: list[ModelMessage] = [call] + result = await cap.compact(messages, _make_ctx()) + assert result[0] is messages[0] + + @pytest.mark.anyio + async def test_request_messages_and_other_parts_untouched(self): + from pydantic_ai.messages import ThinkingPart + + big_user = _user('u' * 5_000) + mixed = ModelResponse(parts=[ThinkingPart(content='t' * 5_000), TextPart(content='z' * 5_000)]) + cap = ClampOversizedMessages(max_part_chars=1_000, keep_head_chars=50, keep_tail_chars=50) + result = await cap.compact([big_user, mixed], _make_ctx()) + + # Request-side message is left as-is (same object). + assert result[0] is big_user + # The thinking part is untouched; only the text part is clamped. + out_mixed = result[1] + assert isinstance(out_mixed, ModelResponse) + thinking, text = out_mixed.parts + assert isinstance(thinking, ThinkingPart) + assert thinking.content == 't' * 5_000 + assert isinstance(text, TextPart) + assert '[clamped: removed' in text.content + + @pytest.mark.anyio + async def test_before_model_request(self): + text = 'q' * 5_000 + cap = ClampOversizedMessages(max_part_chars=1_000, keep_head_chars=50, keep_tail_chars=50) + rc = _make_request_context([_assistant(text)]) + result = await cap.before_model_request(_make_ctx(), rc) + part = result.messages[0].parts[0] + assert isinstance(part, TextPart) + assert '[clamped: removed' in part.content + + def test_marker_format(self): + assert _CLAMP_MARKER.format(removed=10, original=20) == '\n[clamped: removed 10 of 20 characters]\n' + + # --------------------------------------------------------------------------- # Public path — Agent(capabilities=[...]) # --------------------------------------------------------------------------- @@ -1820,6 +1978,18 @@ class TestPublicPath: result = await agent.run('hello') assert result.output is not None + @pytest.mark.anyio + async def test_clamp_oversized_wired_into_agent(self): + from pydantic_ai import Agent + from pydantic_ai.models.test import TestModel + + agent = Agent( + TestModel(), + capabilities=[ClampOversizedMessages(max_part_chars=1)], + ) + result = await agent.run('hello') + assert result.output is not None + # --------------------------------------------------------------------------- # Remaining branch coverage — defensive paths in shared helpers