diff --git a/haystack/components/generators/chat/utils.py b/haystack/components/generators/chat/utils.py new file mode 100644 index 00000000000..9a1a6d5b824 --- /dev/null +++ b/haystack/components/generators/chat/utils.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from haystack.components.generators.chat.types import ChatGenerator + +_CHAT_COMPLETIONS_GENERATORS = {"OpenAIChatGenerator", "AzureOpenAIChatGenerator"} +_RESPONSES_GENERATORS = {"OpenAIResponsesChatGenerator", "AzureOpenAIResponsesChatGenerator"} + + +def _generator_output_token_limit_key(chat_generator: ChatGenerator) -> str | None: + """Return the output-token parameter used by a known built-in Chat Generator.""" + class_names = {cls.__name__ for cls in type(chat_generator).__mro__} + if class_names & _RESPONSES_GENERATORS: + return "max_output_tokens" + if class_names & _CHAT_COMPLETIONS_GENERATORS: + return "max_completion_tokens" + return None + + +def _resolve_output_token_limit(chat_generator: ChatGenerator, default_limit: int) -> tuple[int, dict[str, Any] | None]: + """ + Resolve an effective output-token limit and runtime kwargs for a Chat Generator. + + A recognized limit configured directly on a built-in generator wins and is not repeated at runtime. When the + built-in generator has no configured limit, the default is returned as its provider-specific runtime setting. + Unknown generators receive no runtime setting because the ChatGenerator protocol does not standardize the key. + + :param chat_generator: The generator whose output should be limited. + :param default_limit: The positive fallback output-token limit. + :returns: The effective limit and provider-specific runtime generation kwargs, or None for no runtime kwargs. + """ + limit_key = _generator_output_token_limit_key(chat_generator=chat_generator) + if limit_key is None: + return default_limit, None + + configured = getattr(chat_generator, "generation_kwargs", None) + if isinstance(configured, dict) and limit_key in configured: + value = configured[limit_key] + if isinstance(value, int) and not isinstance(value, bool) and value > 0: + return value, None + # The generator owns this setting. Do not silently replace an invalid value; let it report the problem. + return default_limit, None + return default_limit, {limit_key: default_limit} diff --git a/haystack/hooks/compaction/__init__.py b/haystack/hooks/compaction/__init__.py index 878515a6d1f..fb0ae0a1630 100644 --- a/haystack/hooks/compaction/__init__.py +++ b/haystack/hooks/compaction/__init__.py @@ -10,6 +10,7 @@ _import_structure = { "hooks": ["CompactionHook"], "sliding_window": ["SlidingWindowCompactor"], + "summarization": ["SummarizationCompactor"], "tool_result_pruning": ["ToolResultPruningCompactor"], "types": ["Compactor"], } @@ -17,6 +18,7 @@ if TYPE_CHECKING: from .hooks import CompactionHook as CompactionHook from .sliding_window import SlidingWindowCompactor as SlidingWindowCompactor + from .summarization import SummarizationCompactor as SummarizationCompactor from .tool_result_pruning import ToolResultPruningCompactor as ToolResultPruningCompactor from .types import Compactor as Compactor else: diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index 7f99bb8e0b1..492fb27a895 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -7,7 +7,16 @@ from haystack.core.serialization import default_to_dict from haystack.dataclasses import ChatMessage, ChatRole from haystack.hooks.compaction.types import Compactor -from haystack.hooks.compaction.utils import _COMPACTION_META_KEY, _agent_step_spans +from haystack.hooks.compaction.utils import ( + _COMPACTION_META_KEY, + _current_step_groups, + _historical_turn_groups, + _is_compaction_message, + _latest_user_index, + _leading_system_end, + _messages_at, + _messages_except, +) from haystack.token_counters import TokenCounter from haystack.utils.experimental import _experimental @@ -23,86 +32,9 @@ ) -def _leading_system_end(messages: list[ChatMessage]) -> int: - """Return the end of the leading system-message block.""" - for index, message in enumerate(messages): - # Find the first non-leading system message or a system message produced by compaction - if not message.is_from(role=ChatRole.SYSTEM) or _COMPACTION_META_KEY in message.meta: - return index - return len(messages) - - -def _latest_user_index(messages: list[ChatMessage]) -> int | None: - """ - Return the latest user message not produced by compaction. - - :param messages: The conversation to analyze, oldest to newest. - """ - # We loop backwards to find the latest user message - for index in reversed(range(len(messages))): - message = messages[index] - # Find the latest user message that was not produced by a previous compaction - if message.is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in message.meta: - return index - return None - - def _is_compaction_note(message: ChatMessage) -> bool: """Whether a message is an omission note this strategy left in place of removed history.""" - marker = message.meta.get(_COMPACTION_META_KEY) - return message.is_from(role=ChatRole.USER) and isinstance(marker, dict) and marker.get("strategy") == _STRATEGY - - -def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> list[tuple[int, int]]: - """ - Return spans for complete user turns in a bounded section of conversation history. - - Each turn begins with a real user message and continues up to, but does not include, the next real user message. - This groups a user's request with every assistant step and tool result produced in response to it. - - :param messages: The full conversation to analyze, ordered oldest to newest. - :param start: The inclusive index at which to begin looking for historical turns. - :param end: The exclusive index at which to stop. This is normally the current task's user-message index. - :returns: Ordered `(start_index, end_index)` pairs for each complete historical turn. Both indices refer to - `messages`, and `end_index` is exclusive, so a returned pair can be used directly as `messages[start:end]`. - """ - # Reject any user-role message an earlier compaction produced, whichever strategy made it: none of them are user - # requests, so none of them begin a turn. Leaving them out also lets this compaction fold an old note away. - user_indices = [ - index - for index in range(start, end) - if messages[index].is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in messages[index].meta - ] - - # A real user message closes the preceding turn and starts the next one. The final historical turn extends to the - # supplied boundary, which is typically where the protected current task begins. - return [ - (index, user_indices[position + 1] if position + 1 < len(user_indices) else end) - for position, index in enumerate(user_indices) - ] - - -def _index_groups( - messages: list[ChatMessage], spans: list[tuple[int, int]], skip_compaction_notes: bool = False -) -> list[list[int]]: - """ - Expand each span into the message indices it covers, optionally dropping messages an earlier compaction produced. - """ - return [ - [index for index in range(start, end) if not (skip_compaction_notes and _is_compaction_note(messages[index]))] - for start, end in spans - ] - - -def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: - """Return the messages at the given indices, in the order the indices are given.""" - return [messages[index] for index in indices] - - -def _messages_except(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: - """Return the messages the given indices leave out, in conversation order.""" - left_out = set(indices) - return [message for index, message in enumerate(messages) if index not in left_out] + return _is_compaction_message(message=message, strategy=_STRATEGY, role=ChatRole.USER) def _flatten(groups: list[list[int]]) -> list[int]: @@ -124,19 +56,13 @@ def _removable_groups( kept or removed entire, which is what keeps a tool call with its results and an assistant reply with the user message it answers. """ - # Steps belong to the current task, so they start after its anchor, or after the instructions when the - # conversation has no user message to anchor on. - step_start = (task_index + 1) if task_index is not None else system_end - step_groups = _index_groups(messages=messages, spans=_agent_step_spans(messages=messages, start=step_start)) - # An earlier compaction's note is left out of its turn, so keeping the turn folds that note into the note this # compaction leaves behind. - historical_end = task_index if task_index is not None else system_end - historical_groups = _index_groups( - messages=messages, - spans=_historical_turn_spans(messages=messages, start=system_end, end=historical_end), - skip_compaction_notes=True, - ) + historical_groups = [ + [index for index in group if not _is_compaction_note(message=messages[index])] + for group in _historical_turn_groups(messages=messages, system_end=system_end, task_index=task_index) + ] + step_groups = _current_step_groups(messages=messages, system_end=system_end, task_index=task_index) return historical_groups, step_groups diff --git a/haystack/hooks/compaction/summarization.py b/haystack/hooks/compaction/summarization.py new file mode 100644 index 00000000000..9bef7c06abf --- /dev/null +++ b/haystack/hooks/compaction/summarization.py @@ -0,0 +1,503 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from haystack import logging +from haystack.components.generators.chat.types import ChatGenerator +from haystack.components.generators.chat.utils import _resolve_output_token_limit +from haystack.core.serialization import component_to_dict, default_from_dict, default_to_dict +from haystack.dataclasses import ChatMessage, FileContent, ImageContent +from haystack.dataclasses.chat_message import ChatMessageContentT +from haystack.hooks.compaction.types import Compactor +from haystack.hooks.compaction.utils import ( + _COMPACTION_META_KEY, + _current_step_groups, + _historical_turn_groups, + _is_compaction_message, + _latest_user_index, + _leading_system_end, + _messages_at, + _messages_except, +) +from haystack.token_counters import TokenCounter +from haystack.token_counters.utils import _rendered_conversation +from haystack.utils.async_utils import _execute_component_async +from haystack.utils.deserialization import deserialize_component_inplace +from haystack.utils.experimental import _experimental + +logger = logging.getLogger(__name__) + +# Recorded as the strategy on every summary this compactor produces, so a later run can recognize its own summaries. +_STRATEGY = "summarization" + +# Recorded as the `source` on a summary, naming the stretch of conversation it stands in for. Compaction gives these up +# in order, so the Agent's current task is the last thing to go. +_HISTORICAL_TURNS = "historical_turns" +_HISTORICAL_SUMMARIES = "historical_summaries" +_CURRENT_TASK_SUMMARIES = "current_task_summaries" +_CURRENT_TASK_STEPS = "current_task_steps" + +_DEFAULT_SUMMARY_INSTRUCTION = """You are compacting one portion of a conversation between a user and an AI agent so \ +the agent can keep working with fewer tokens. You are shown only the portion being replaced. The rest of the \ +conversation, including the user's current request, stays in place and is not shown to you. Summarize only what you \ +are given, and never say or imply that something did not happen just because it is absent from this portion. + +Use these sections, in this order. Keep every section, and write "(none)" when this portion says nothing about it. + +## Objective +What the user was trying to accomplish, if this portion shows it. + +## Decisions and constraints +Choices made and the reasoning behind them, and any requirements, preferences, or instructions the user gave. Note \ +options that were rejected and why. + +## Work completed +What was done, and what the tool results established. + +## Identifiers +Exact file paths, URLs, IDs, names, commands, and error strings, copied character for character. Images and files \ +appear only as and placeholders; their contents are not available to you and are lost once \ +this portion is replaced, so copy the placeholder details here. + +## Unresolved +Work still outstanding, and the immediate next step. + +Rules: +- Record only what this portion shows. Do not infer, do not give advice, and do not add anything that is not here. +- Copy identifiers exactly rather than describing them. They cannot be recovered once this portion is gone. +- Fold any blocks you are given into your own: keep what is still true, drop what is now \ +stale, and merge in the new facts. +- Use terse bullets. Do not address the user, and do not mention that you are summarizing.""" + + +def _identifying_details(metadata: dict[str, Any]) -> list[str]: + """Render an attachment's metadata as `key=value` pairs, such as the path a file was loaded from.""" + # For ease, we don't support nested keys, we are mostly interested in the top-level keys that identify the + # attachment, such as a file path or URL. + return [f"{key}={value}" for key, value in sorted(metadata.items()) if isinstance(value, (str, int, float, bool))] + + +def _attachment_placeholder(content: ChatMessageContentT) -> str: + """ + Render a placeholder for an attachment that cannot survive summarization, so the summary can preserve its identity. + """ + if isinstance(content, ImageContent): + # Images have no filename, so whatever identifies one lives in its `meta`. + return f"" + if isinstance(content, FileContent): + details = [content.filename or "unnamed", content.mime_type or "unknown type"] + return f"" + return f"<{type(content).__name__}>" + + +def _is_summary(message: ChatMessage) -> bool: + """Whether a message is a summary this strategy wrote.""" + return _is_compaction_message(message=message, strategy=_STRATEGY) + + +def _previous_summary_indices(messages: list[ChatMessage], start: int, end: int) -> list[int]: + """Return the positions of the summaries an earlier compaction left in a bounded part of a conversation.""" + return [index for index in range(start, end) if _is_summary(message=messages[index])] + + +def _raw_historical_turn_groups( + messages: list[ChatMessage], system_end: int, task_index: int | None +) -> list[list[int]]: + """ + Return the historical turns that still hold raw, never-summarized conversation, oldest turn first. + + Summaries an earlier compaction wrote are excluded, so summarizing a turn leaves them in place for + `_HISTORICAL_SUMMARIES` to fold later. The list is empty when there are no historical turns, or when every one of + them is already nothing but summaries. + """ + # Strip the previous summaries out of each turn, then drop the turns that strip away to nothing. + groups = [ + [index for index in group if not _is_summary(message=messages[index])] + for group in _historical_turn_groups(messages=messages, system_end=system_end, task_index=task_index) + ] + return [group for group in groups if group] + + +def _groups_to_summarize( + messages: list[ChatMessage], + groups: list[list[int]], + target_tokens: int, + summary_tokens: int, + token_counter: TokenCounter, +) -> list[int]: + """ + Return the fewest oldest groups whose removal makes room for a summary of the expected size. + + Groups are taken oldest first and counting stops as soon as what remains, plus the summary that replaces them, + fits the target. When even taking all of them is not enough, all of them are returned. + """ + selected: list[int] = [] + for group in groups: + selected.extend(group) + remaining = token_counter.count(messages=_messages_except(messages=messages, indices=selected)) + if remaining + summary_tokens <= target_tokens: + break + return selected + + +def _summary_message(text: str, summarized_messages: int, source: str) -> ChatMessage: + """Build the marked user message that stands in for the messages the summary replaced.""" + body = f"\n{text.strip()}\n" + meta = {_COMPACTION_META_KEY: {"strategy": _STRATEGY, "summarized_messages": summarized_messages, "source": source}} + return ChatMessage.from_user(text=body, meta=meta) + + +def _replace_indices(messages: list[ChatMessage], indices: list[int], summary: ChatMessage) -> list[ChatMessage]: + """Replace the selected messages, which need not be contiguous, with one summary at the oldest one's position.""" + selected = set(indices) + # The summary stands in for everything it replaced, so it takes the position of the oldest message it covers. + insertion_index = min(indices) + compacted: list[ChatMessage] = [] + for index, message in enumerate(messages): + # Emit the summary before the message it displaces, so the surrounding conversation keeps its order. + if index == insertion_index: + compacted.append(summary) + if index not in selected: + compacted.append(message) + return compacted + + +@_experimental +class SummarizationCompactor(Compactor): + """ + Condenses old historical turns first, then old steps from the Agent's current task. + + Leading system messages and the latest real user message are always kept. Historical turns are summarized in full, + oldest first. Summaries normally accumulate so they are not repeatedly rewritten; if every historical turn has + already been summarized and more space is needed, those historical summaries are folded into one before any + current-task steps are summarized. An assistant message and all immediately following tool results form one step, + so tool calls are never separated from their results. + + Each summary is requested within `max_summary_tokens`. Built-in OpenAI Chat Completions and Responses generators + receive the corresponding runtime output limit unless they already configure one in their `generation_kwargs`, in + which case the generator's setting wins. Other generators receive the same limit as prompt guidance and the actual + result is measured before it is accepted. + + ```python + from haystack.components.agents import Agent + from haystack.components.generators.chat import OpenAIResponsesChatGenerator + from haystack.hooks.compaction import CompactionHook, SummarizationCompactor + + summary_generator = OpenAIResponsesChatGenerator(model="gpt-5.4-nano") + hook = CompactionHook( + compactor=SummarizationCompactor(chat_generator=summary_generator), + context_window=400_000, + compact_at=0.7, + compact_to=0.4, + ) + agent = Agent(chat_generator=agent_generator, tools=[web_search], hooks={"before_llm": [hook]}) + ``` + """ + + def __init__( + self, + chat_generator: ChatGenerator, + *, + min_keep_steps: int = 1, + max_summary_tokens: int = 1024, + summary_instruction: str = _DEFAULT_SUMMARY_INSTRUCTION, + raise_on_failure: bool = False, + ) -> None: + """ + Initialize the compactor. + + :param chat_generator: The Chat Generator used to write summaries. Its configured output-token limit takes + precedence over `max_summary_tokens` when the generator exposes a recognized setting. + :param min_keep_steps: The fewest complete recent Agent steps to keep, even when they exceed the target. + :param max_summary_tokens: The output-token budget reserved for each summary. Known built-in generators receive + the corresponding runtime generation setting unless one is already configured on the generator. + :param summary_instruction: What the model is told to preserve when it writes a summary. The default asks for + fixed sections covering the objective, decisions and constraints, completed work, exact identifiers, and + unresolved work, each written as `(none)` when the summarized portion says nothing about it. It also states + that only part of the conversation is shown, so the model does not conclude that something never happened + just because it is absent. The token budget is appended to whatever is given here, so a replacement does + not need to mention it. + :param raise_on_failure: Whether a failed or non-shrinking summarization raises. By default the failure is + logged and any successful partial compaction is returned. + :raises ValueError: If `min_keep_steps` is negative or `max_summary_tokens` is not positive. + """ + if min_keep_steps < 0: + raise ValueError(f"`min_keep_steps` must be at least 0, got {min_keep_steps}.") + if max_summary_tokens < 1: + raise ValueError(f"`max_summary_tokens` must be a positive number of tokens, got {max_summary_tokens}.") + self.chat_generator = chat_generator + self.min_keep_steps = min_keep_steps + self.max_summary_tokens = max_summary_tokens + self.summary_instruction = summary_instruction + self.raise_on_failure = raise_on_failure + + def compact( + self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter + ) -> list[ChatMessage] | None: + """ + Return a progressively summarized conversation, or None when no useful reduction is possible. + + :param messages: The conversation to compact, ordered oldest to newest. + :param target_tokens: The token budget the compacted messages should aim to fit. + :param token_counter: The counter used both to plan compaction and verify generated summaries. + :returns: A smaller replacement conversation, or None when nothing was reduced. + """ + # How large each summary may be. A recognized built-in generator is held to it by its own runtime setting; + # any other generator gets it as prompt guidance only, since the protocol guarantees nothing beyond `run`. + summary_tokens, generation_kwargs = _resolve_output_token_limit( + chat_generator=self.chat_generator, default_limit=self.max_summary_tokens + ) + run_kwargs = {"generation_kwargs": generation_kwargs} if generation_kwargs else {} + + # Rebound only when a summary is applied, and never mutated, so `messages` is left as the caller passed it. + compacted = messages + summarized = False + while True: + # Ask which stretch of the conversation to give up next. None means the target is met or nothing is left. + plan = self._next_summary( + messages=compacted, + target_tokens=target_tokens, + token_counter=token_counter, + summary_tokens=summary_tokens, + ) + if plan is None: + break + indices, source = plan + prompt = self._prompt(messages=compacted, indices=indices, summary_tokens=summary_tokens) + try: + # Summarize that stretch and swap it in, so the next round plans against the smaller conversation. + # A generator error or a summary that does not shrink raises out of here. + result = self.chat_generator.run(messages=prompt, **run_kwargs) + compacted = self._apply_summary( + messages=compacted, indices=indices, source=source, result=result, token_counter=token_counter + ) + summarized = True + except Exception as error: + # Stop at the last summary that worked, unless `raise_on_failure` says to propagate. + self._report_failure(error=error) + break + # Every applied summary was measured as shrinking the conversation, so any summary at all is real progress, + # whether or not the target was met. Without one there is nothing to hand back. + return compacted if summarized else None + + async def compact_async( + self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter + ) -> list[ChatMessage] | None: + """ + Asynchronously return a progressively summarized conversation. + + :param messages: The conversation to compact, ordered oldest to newest. + :param target_tokens: The token budget the compacted messages should aim to fit. + :param token_counter: The counter used both to plan compaction and verify generated summaries. + :returns: A smaller replacement conversation, or None when nothing was reduced. + """ + # How large each summary may be. A recognized built-in generator is held to it by its own runtime setting; + # any other generator gets it as prompt guidance only, since the protocol guarantees nothing beyond `run`. + summary_tokens, generation_kwargs = _resolve_output_token_limit( + chat_generator=self.chat_generator, default_limit=self.max_summary_tokens + ) + run_kwargs = {"generation_kwargs": generation_kwargs} if generation_kwargs else {} + + # Rebound only when a summary is applied, and never mutated, so `messages` is left as the caller passed it. + compacted = messages + summarized = False + while True: + # Ask which stretch of the conversation to give up next. None means the target is met or nothing is left. + plan = self._next_summary( + messages=compacted, + target_tokens=target_tokens, + token_counter=token_counter, + summary_tokens=summary_tokens, + ) + if plan is None: + break + indices, source = plan + prompt = self._prompt(messages=compacted, indices=indices, summary_tokens=summary_tokens) + try: + # Summarize that stretch and swap it in, so the next round plans against the smaller conversation. + # Only the generator call is awaited; planning and swapping are pure. + result = await _execute_component_async( + component_instance=self.chat_generator, messages=prompt, **run_kwargs + ) + compacted = self._apply_summary( + messages=compacted, indices=indices, source=source, result=result, token_counter=token_counter + ) + summarized = True + except Exception as error: + # Stop at the last summary that worked, unless `raise_on_failure` says to propagate. + self._report_failure(error=error) + break + # Every applied summary was measured as shrinking the conversation, so any summary at all is real progress, + # whether or not the target was met. Without one there is nothing to hand back. + return compacted if summarized else None + + def _next_summary( + self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, summary_tokens: int + ) -> tuple[list[int], str] | None: + """ + Choose the next stretch of conversation to replace with a summary. + + Four tiers are tried in order, so the oldest and least useful context goes first and the Agent's current task + is given up last: + + 1. `_HISTORICAL_TURNS`: the fewest oldest raw turns that make room for a summary. + 2. `_HISTORICAL_SUMMARIES`: nothing raw is left in history, so fold its summaries into one. + 3. `_CURRENT_TASK_SUMMARIES`: fold the summaries earlier steps left behind before giving up more steps. + 4. `_CURRENT_TASK_STEPS`: the fewest oldest steps of the current task, keeping `min_keep_steps` of the newest. + + :param messages: The conversation as it stands, ordered oldest to newest. + :param target_tokens: The token budget the conversation should come in under. + :param token_counter: The counter used to measure candidate selections. + :param summary_tokens: The size a summary is expected to take, reserved when choosing how much to replace. + :returns: The message indices to summarize and the `source` to record on the resulting summary, or None when + the conversation already fits or nothing is left that may be given up. + """ + # Nothing to give up once the conversation fits. + if token_counter.count(messages=messages) <= target_tokens: + return None + + # The landmarks everything is measured against: the Agent's instructions, and the user message anchoring the + # current task. History runs from the instructions up to that anchor, the current task from the anchor on. + system_end = _leading_system_end(messages=messages) + task_index = _latest_user_index(messages=messages) + history_end = task_index if task_index is not None else system_end + task_start = task_index + 1 if task_index is not None else system_end + + # Tier 1. Raw history is the cheapest context to lose, so take the oldest turns that still hold any. + historical_turns = _raw_historical_turn_groups(messages=messages, system_end=system_end, task_index=task_index) + if historical_turns: + oldest_turns = _groups_to_summarize( + messages=messages, + groups=historical_turns, + target_tokens=target_tokens, + summary_tokens=summary_tokens, + token_counter=token_counter, + ) + return oldest_turns, _HISTORICAL_TURNS + + # Tier 2. History is nothing but summaries now, so the only room left there is in folding them into one. They + # are left to accumulate until this point so that they are not rewritten on every compaction. + history_summaries = _previous_summary_indices(messages=messages, start=system_end, end=history_end) + if len(history_summaries) > 1: + return history_summaries, _HISTORICAL_SUMMARIES + + # History is exhausted, so the current task has to pay. Its `min_keep_steps` newest steps are off limits. + agent_steps = _current_step_groups(messages=messages, system_end=system_end, task_index=task_index) + eligible_steps = agent_steps[: max(len(agent_steps) - self.min_keep_steps, 0)] + if not eligible_steps: + return None + + # Tier 3. Fold the summaries earlier steps left behind before spending another raw step on the same space. + task_summaries = _previous_summary_indices(messages=messages, start=task_start, end=len(messages)) + if len(task_summaries) > 1: + return task_summaries, _CURRENT_TASK_SUMMARIES + + # Tier 4. Last resort: give up the oldest steps of the task the Agent is working on right now. + oldest_steps = _groups_to_summarize( + messages=messages, + groups=eligible_steps, + target_tokens=target_tokens, + summary_tokens=summary_tokens, + token_counter=token_counter, + ) + return oldest_steps, _CURRENT_TASK_STEPS + + def _prompt(self, messages: list[ChatMessage], indices: list[int], summary_tokens: int) -> list[ChatMessage]: + """Build the bounded summarization instruction and the rendered transcript of the selected messages.""" + transcript = _rendered_conversation( + _messages_at(messages=messages, indices=indices), placeholder=_attachment_placeholder + ) + instruction = ( + f"{self.summary_instruction}\n\nWrite a complete summary in no more than approximately " + f"{summary_tokens} tokens. Prioritize completeness within that limit so the response is not cut off." + ) + return [ + ChatMessage.from_system(text=instruction), + ChatMessage.from_user(text=f"\n{transcript}\n"), + ] + + @staticmethod + def _apply_summary( + messages: list[ChatMessage], + indices: list[int], + source: str, + result: dict[str, Any], + token_counter: TokenCounter, + ) -> list[ChatMessage]: + """ + Swap the selected messages for the generated summary. + + :raises RuntimeError: If the generator returned no usable text, or if the swap did not make the conversation + smaller, in which case keeping the raw messages is the better outcome. + """ + replies = result.get("replies") or [] + text = replies[-1].text if replies else None + if not text or not text.strip(): + raise RuntimeError("The Chat Generator returned no text to use as a conversation summary.") + summary = _summary_message(text=text, summarized_messages=len(indices), source=source) + compacted = _replace_indices(messages=messages, indices=indices, summary=summary) + before = token_counter.count(messages=messages) + after = token_counter.count(messages=compacted) + if after >= before: + raise RuntimeError( + f"The generated summary did not reduce the conversation size ({before} tokens before and {after} " + "tokens after)." + ) + return compacted + + def _report_failure(self, error: Exception) -> None: + """Re-raise a failed summarization or log it, so whatever compacted successfully so far is still returned.""" + if self.raise_on_failure: + raise error + logger.warning( + "Summarizing the conversation for context compaction failed; keeping the last successful result. " + "Error: {error}", + error=error, + ) + + def warm_up(self) -> None: + """Warm up the Chat Generator that writes summaries.""" + if hasattr(self.chat_generator, "warm_up"): + self.chat_generator.warm_up() + + async def warm_up_async(self) -> None: + """Warm up the Chat Generator on the serving event loop.""" + warm_up_async = getattr(self.chat_generator, "warm_up_async", None) + if warm_up_async is not None: + await warm_up_async() + elif hasattr(self.chat_generator, "warm_up"): + self.chat_generator.warm_up() + + def close(self) -> None: + """Release the Chat Generator's resources.""" + if hasattr(self.chat_generator, "close"): + self.chat_generator.close() + + async def close_async(self) -> None: + """Release the Chat Generator's resources.""" + close_async = getattr(self.chat_generator, "close_async", None) + if close_async is not None: + await close_async() + elif hasattr(self.chat_generator, "close"): + self.chat_generator.close() + + def to_dict(self) -> dict[str, Any]: + """Serialize the compactor and its Chat Generator.""" + return default_to_dict( + self, + chat_generator=component_to_dict(obj=self.chat_generator, name="chat_generator"), + min_keep_steps=self.min_keep_steps, + max_summary_tokens=self.max_summary_tokens, + summary_instruction=self.summary_instruction, + raise_on_failure=self.raise_on_failure, + ) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "SummarizationCompactor": + """Deserialize the compactor and reconstruct its Chat Generator.""" + init_params = data.get("init_parameters", {}) + if init_params.get("chat_generator") is not None: + deserialize_component_inplace(data=init_params, key="chat_generator") + return default_from_dict(cls=cls, data=data) diff --git a/haystack/hooks/compaction/utils.py b/haystack/hooks/compaction/utils.py index 6a6d1402c52..590ddb84b2d 100644 --- a/haystack/hooks/compaction/utils.py +++ b/haystack/hooks/compaction/utils.py @@ -10,6 +10,41 @@ _COMPACTION_META_KEY = "context_compaction" +def _leading_system_end(messages: list[ChatMessage]) -> int: + """Return the end of the leading system-message block, excluding system messages created by compaction.""" + for index, message in enumerate(messages): + if not message.is_from(role=ChatRole.SYSTEM) or _COMPACTION_META_KEY in message.meta: + return index + return len(messages) + + +def _latest_user_index(messages: list[ChatMessage]) -> int | None: + """Return the latest user message not produced by compaction.""" + for index in reversed(range(len(messages))): + message = messages[index] + if message.is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in message.meta: + return index + return None + + +def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: + """Return the messages at the given indices, in the order the indices are given.""" + return [messages[index] for index in indices] + + +def _messages_except(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: + """Return the messages the given indices leave out, in conversation order.""" + left_out = set(indices) + return [message for index, message in enumerate(messages) if index not in left_out] + + +def _is_compaction_message(message: ChatMessage, strategy: str, role: ChatRole | None = None) -> bool: + """Return whether a message was produced by a compaction strategy and optionally has the requested role.""" + marker = message.meta.get(_COMPACTION_META_KEY) + has_role = role is None or message.is_from(role=role) + return has_role and isinstance(marker, dict) and marker.get("strategy") == strategy + + def _last_assistant_index(messages: list[ChatMessage]) -> int: """Return the index of the last assistant message, or -1 if none exists.""" for index in reversed(range(len(messages))): @@ -47,6 +82,38 @@ def _agent_step_spans(messages: list[ChatMessage], start: int) -> list[tuple[int return spans +def _current_step_groups(messages: list[ChatMessage], system_end: int, task_index: int | None) -> list[list[int]]: + """Return message-index groups for complete Agent steps belonging to the current task.""" + step_start = task_index + 1 if task_index is not None else system_end + return [list(range(start, end)) for start, end in _agent_step_spans(messages=messages, start=step_start)] + + +def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> list[tuple[int, int]]: + """ + Return spans for complete real-user turns in a bounded section of conversation history. + + Each turn begins with a user message not created by compaction and continues up to the next such message. + """ + user_indices = [ + index + for index in range(start, end) + if messages[index].is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in messages[index].meta + ] + return [ + (index, user_indices[position + 1] if position + 1 < len(user_indices) else end) + for position, index in enumerate(user_indices) + ] + + +def _historical_turn_groups(messages: list[ChatMessage], system_end: int, task_index: int | None) -> list[list[int]]: + """Return message-index groups for complete historical turns preceding the current task.""" + historical_end = task_index if task_index is not None else system_end + return [ + list(range(start, end)) + for start, end in _historical_turn_spans(messages=messages, start=system_end, end=historical_end) + ] + + def _estimated_context_tokens( messages: list[ChatMessage], context_tokens: int, token_counter: TokenCounter, tools: ToolsType | None = None ) -> int: diff --git a/haystack/token_counters/utils.py b/haystack/token_counters/utils.py index 9dbf9b77940..7f4d8f3ef71 100644 --- a/haystack/token_counters/utils.py +++ b/haystack/token_counters/utils.py @@ -3,11 +3,15 @@ # SPDX-License-Identifier: Apache-2.0 import json +from collections.abc import Callable from haystack.dataclasses import ChatMessage, FileContent, ImageContent, TextContent from haystack.dataclasses.chat_message import ChatMessageContentT, ToolCallResultContentT from haystack.tools import ToolsType, flatten_tools_or_toolsets +# Builds the stand-in for message content that has no text form, such as an image. +_PlaceholderFn = Callable[[ChatMessageContentT], str] + def _non_text_placeholder(content: ChatMessageContentT) -> str: """A short stand-in, such as ``, for message content that has no text form.""" @@ -18,26 +22,33 @@ def _non_text_placeholder(content: ChatMessageContentT) -> str: return f"<{type(content).__name__}>" -def _tool_result_text(result: ToolCallResultContentT) -> str: +def _tool_result_text(result: ToolCallResultContentT, placeholder: _PlaceholderFn = _non_text_placeholder) -> str: """A tool result as a single string, with placeholders standing in for any non-text parts.""" if isinstance(result, str): return result - return "".join(block.text if isinstance(block, TextContent) else _non_text_placeholder(block) for block in result) + return "".join(block.text if isinstance(block, TextContent) else placeholder(block) for block in result) -def _render_message(message: ChatMessage) -> str: +def _render_message(message: ChatMessage, placeholder: _PlaceholderFn = _non_text_placeholder) -> str: """ One message as one or more lines of plain text. Reasoning content is deliberately left out: providers discard it between turns, so it is not part of the context being measured. + + :param message: The message to render. + :param placeholder: Builds the stand-in for content that has no text form. The default is short and stable, which + is what a counter needs because the stand-in's own length is what gets measured. A caller rendering for a model + to read can pass one that describes the content instead. + :returns: The rendered message. """ role = message.role.value # A tool-result msg only carries tool_call_results, so it is rendered on its own and labelled with the tool that # produced it. if results := message.tool_call_results: return "\n".join( - f"[tool:{result.origin.tool_name}{' (error)' if result.error else ''}] {_tool_result_text(result.result)}" + f"[tool:{result.origin.tool_name}{' (error)' if result.error else ''}] " + f"{_tool_result_text(result.result, placeholder=placeholder)}" for result in results ) @@ -50,13 +61,13 @@ def _render_message(message: ChatMessage) -> str: # Images and files cost tokens too, so they need a stand-in rather than being skipped. non_text: list[ChatMessageContentT] = [*message.images, *message.files] for content in non_text: - lines.append(f"[{role}] {_non_text_placeholder(content)}") + lines.append(f"[{role}] {placeholder(content)}") return "\n".join(lines) if lines else f"[{role}] " -def _rendered_conversation(messages: list[ChatMessage]) -> str: +def _rendered_conversation(messages: list[ChatMessage], *, placeholder: _PlaceholderFn = _non_text_placeholder) -> str: """The whole conversation as one plain-text block, which is what a counter measures.""" - return "\n".join(_render_message(message) for message in messages) + return "\n".join(_render_message(message, placeholder=placeholder) for message in messages) def _rendered_tools(tools: ToolsType | None) -> str: diff --git a/pydoc/hooks_api.yml b/pydoc/hooks_api.yml index 4fed6fb1452..fef10aad17e 100644 --- a/pydoc/hooks_api.yml +++ b/pydoc/hooks_api.yml @@ -1,7 +1,7 @@ loaders: - search_path: [../haystack/hooks] modules: ["protocol", "from_function", "compaction/hooks", "compaction/sliding_window", - "compaction/tool_result_pruning", + "compaction/summarization", "compaction/tool_result_pruning", "compaction/types/protocol", "human_in_the_loop/dataclasses", "human_in_the_loop/hooks", "human_in_the_loop/policies", "human_in_the_loop/strategies", "human_in_the_loop/user_interfaces", "tool_result_offloading/hooks", "tool_result_offloading/policies", "tool_result_offloading/stores", diff --git a/releasenotes/notes/add-summarization-compactor-91b6be6855f478df.yaml b/releasenotes/notes/add-summarization-compactor-91b6be6855f478df.yaml new file mode 100644 index 00000000000..d6125f531b5 --- /dev/null +++ b/releasenotes/notes/add-summarization-compactor-91b6be6855f478df.yaml @@ -0,0 +1,6 @@ +--- +features: + - | + Added the experimental ``SummarizationCompactor``. It uses a Chat Generator to condense the Agent's oldest + historical turns, accumulated summaries, and old current-task steps as needed to reach the context target while + preserving leading system messages, the latest user task, and complete recent tool-calling steps. diff --git a/test/components/generators/chat/test_utils.py b/test/components/generators/chat/test_utils.py new file mode 100644 index 00000000000..6d64d54dd26 --- /dev/null +++ b/test/components/generators/chat/test_utils.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from haystack.components.generators.chat import MockChatGenerator +from haystack.components.generators.chat.utils import _generator_output_token_limit_key, _resolve_output_token_limit + + +class OpenAIChatGenerator(MockChatGenerator): + def __init__(self, generation_kwargs=None): + super().__init__("response") + self.generation_kwargs = generation_kwargs or {} + + +class OpenAIResponsesChatGenerator(MockChatGenerator): + def __init__(self, generation_kwargs=None): + super().__init__("response") + self.generation_kwargs = generation_kwargs or {} + + +class CustomGenerator(MockChatGenerator): + generation_kwargs = {"max_tokens": 7} + + +def test_resolves_chat_completions_limit(): + generator = OpenAIChatGenerator() + assert _generator_output_token_limit_key(generator) == "max_completion_tokens" + assert _resolve_output_token_limit(generator, 100) == (100, {"max_completion_tokens": 100}) + + +def test_resolves_responses_limit(): + generator = OpenAIResponsesChatGenerator() + assert _generator_output_token_limit_key(generator) == "max_output_tokens" + assert _resolve_output_token_limit(generator, 100) == (100, {"max_output_tokens": 100}) + + +def test_configured_generator_limit_wins_without_mutation(): + generator = OpenAIChatGenerator({"temperature": 0, "max_completion_tokens": 23}) + original = dict(generator.generation_kwargs) + assert _resolve_output_token_limit(generator, 100) == (23, None) + assert generator.generation_kwargs == original + + +def test_unknown_generator_receives_no_guessed_runtime_setting(): + assert _generator_output_token_limit_key(CustomGenerator()) is None + assert _resolve_output_token_limit(CustomGenerator(), 100) == (100, None) diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index 7757b59e7f1..d1186578880 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -6,8 +6,8 @@ from haystack.dataclasses import ChatMessage, ChatRole, ToolCall from haystack.hooks.compaction import SlidingWindowCompactor -from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE, _historical_turn_spans, _is_compaction_note -from haystack.hooks.compaction.utils import _COMPACTION_META_KEY +from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE, _is_compaction_note +from haystack.hooks.compaction.utils import _COMPACTION_META_KEY, _historical_turn_spans from test.hooks.compaction.helpers import ( FakeCounter, count_markers, diff --git a/test/hooks/compaction/test_summarization.py b/test/hooks/compaction/test_summarization.py new file mode 100644 index 00000000000..b2b5197125e --- /dev/null +++ b/test/hooks/compaction/test_summarization.py @@ -0,0 +1,338 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from haystack.components.generators.chat import MockChatGenerator +from haystack.dataclasses import ChatMessage, ChatRole, FileContent, ImageContent, TextContent, ToolCall +from haystack.hooks.compaction import SummarizationCompactor +from haystack.hooks.compaction.summarization import _attachment_placeholder +from haystack.hooks.compaction.utils import _COMPACTION_META_KEY +from test.hooks.compaction.helpers import FakeCounter, tool_call, tool_result + +pytestmark = pytest.mark.filterwarnings("ignore::haystack.utils.experimental.ExperimentalWarning") + +# A target of one token forces every tier to run, isolating the structural rules from sizing. +SMALLEST = 1 +# One character per token, so the padded messages below are obviously the expensive ones. +COUNTER = FakeCounter(chars_per_token=1) + + +def summarizer(*responses: str | Exception) -> tuple[MockChatGenerator, list[str]]: + """ + A Chat Generator returning the given summaries in order, recording the prompt it received for each. + + An `Exception` among the responses is raised instead of answering, so a test can fail one summarization step. + """ + queued = list(responses) + prompts: list[str] = [] + + def respond(messages: list[ChatMessage]) -> str: + prompts.append("\n".join(message.text or "" for message in messages)) + response = queued.pop(0) + if isinstance(response, Exception): + raise response + return response + + return MockChatGenerator(response_fn=respond), prompts + + +def summary(text: str, source: str) -> ChatMessage: + """A summary an earlier compaction left behind, marked the way this compactor marks its own.""" + return ChatMessage.from_user( + f"\n{text}\n", + meta={_COMPACTION_META_KEY: {"strategy": "summarization", "source": source}}, + ) + + +def sources(messages: list[ChatMessage]) -> list[str]: + """Which stretch of conversation each summary in `messages` stands in for, oldest first.""" + return [ + message.meta[_COMPACTION_META_KEY]["source"] for message in messages if _COMPACTION_META_KEY in message.meta + ] + + +def two_turns_and_a_task() -> list[ChatMessage]: + """A padded oldest turn, a short recent turn, and the current task with one step behind it.""" + return [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("oldest question " * 30), + ChatMessage.from_assistant("oldest answer " * 30), + ChatMessage.from_user("recent question"), + ChatMessage.from_assistant("recent answer"), + ChatMessage.from_user("current task"), + ChatMessage.from_assistant("current step"), + ] + + +def a_task_with_two_steps() -> list[ChatMessage]: + """The current task with a padded oldest step and a cheap newest one, and no history in front of it.""" + return [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("current task"), + tool_call("old"), + tool_result("old result " * 30, call_id="old"), + tool_call("new"), + tool_result("new result", call_id="new"), + ] + + +class TestAttachmentPlaceholder: + @pytest.mark.parametrize( + ("content", "expected"), + [ + pytest.param(ImageContent(base64_image="Zm9v", mime_type="image/png"), "", id="image"), + pytest.param( + ImageContent(base64_image="Zm9v", mime_type="image/png", meta={"file_path": "/tmp/shot.png"}), + "", + id="image-named-by-meta", + ), + pytest.param( + FileContent(base64_data="Zm9v", mime_type="application/pdf", filename="q3.pdf"), + "", + id="file", + ), + pytest.param( + FileContent(base64_data="Zm9v", mime_type="application/pdf", extra={"page": 4}), + "", + id="file-unnamed-with-extra", + ), + # A nested value could be arbitrarily large, so it is left out rather than bloating the prompt. + pytest.param( + ImageContent(base64_image="Zm9v", mime_type="image/png", meta={"boxes": [[1, 2], [3, 4]]}), + "", + id="nested-metadata-left-out", + ), + ], + ) + def test_names_the_attachment(self, content, expected): + assert _attachment_placeholder(content) == expected + + +class TestSummarizationCompactor: + def test_summarizes_oldest_historical_turn(self): + messages = two_turns_and_a_task() + generator, prompts = summarizer("short historical summary") + # Room for everything but the padded oldest turn, plus the summary standing in for it. + target_tokens = COUNTER.count([messages[0], *messages[3:]]) + 100 + + compacted = SummarizationCompactor(generator, max_summary_tokens=100).compact( + messages=messages, target_tokens=target_tokens, token_counter=COUNTER + ) + + assert compacted is not None + # Only the oldest turn was summarized, so the recent turn never reached the generator. + assert len(prompts) == 1 + assert "oldest question" in prompts[0] + assert "recent question" not in prompts[0] + assert compacted == [messages[0], compacted[1], *messages[3:]] + + def test_leaves_the_input_conversation_untouched(self): + messages = two_turns_and_a_task() + SummarizationCompactor(MockChatGenerator("summary"), max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + assert messages == two_turns_and_a_task() + + def test_summarizes_historical_turns_then_current_steps(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("old question " * 40), + ChatMessage.from_assistant("old answer " * 40), + *a_task_with_two_steps()[1:], + ] + generator, prompts = summarizer("history", "old step") + + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + + assert compacted is not None + assert "old question" in prompts[0] + assert "old result" in prompts[1] + assert sources(compacted) == ["historical_turns", "current_task_steps"] + # The newest step is never given up. + assert compacted[-2:] == messages[-2:] + + def test_folds_historical_summaries_before_current_steps(self): + messages = [ + ChatMessage.from_system("rules"), + summary("first history " * 20, "historical_turns"), + summary("second history " * 20, "historical_turns"), + *a_task_with_two_steps()[1:], + ] + generator, prompts = summarizer("combined history", "old step") + + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + + assert compacted is not None + # The two historical summaries are folded into one before any current-task step is touched. + assert "first history" in prompts[0] and "old result" not in prompts[0] + assert "old result" in prompts[1] + assert sources(compacted) == ["historical_summaries", "current_task_steps"] + + def test_folds_current_task_summaries_before_more_steps(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("current task"), + summary("first step summary " * 20, "current_task_steps"), + summary("second step summary " * 20, "current_task_steps"), + *a_task_with_two_steps()[2:], + ] + generator, prompts = summarizer("combined steps", "old step") + + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + + assert compacted is not None + assert "first step summary" in prompts[0] + assert "old result" in prompts[1] + assert sources(compacted) == ["current_task_summaries", "current_task_steps"] + assert compacted[-2:] == messages[-2:] + + @pytest.mark.parametrize(("min_keep_steps", "expected"), [(0, 0), (1, 1), (2, 2), (20, 2)]) + def test_min_keep_steps_wins_over_an_unaffordable_target(self, min_keep_steps, expected): + messages = a_task_with_two_steps() + generator, _ = summarizer("step summary") + compacted = SummarizationCompactor(generator, min_keep_steps=min_keep_steps, max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + result = compacted or messages + assert sum(message.is_from(role=ChatRole.ASSISTANT) for message in result) == expected + + def test_attachments_are_named_in_the_transcript(self): + image = ImageContent(base64_image="Zm9v", mime_type="image/png", meta={"file_path": "/tmp/shot.png"}) + pdf = FileContent(base64_data="Zm9v", mime_type="application/pdf", filename="q3.pdf") + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user(content_parts=["review this " * 20, pdf]), + tool_call("c1"), + # An attachment a tool returned is nested inside the tool result rather than on the message. + ChatMessage.from_tool( + tool_result=[TextContent(text="captured " * 20), image], + origin=ToolCall(tool_name="browse", arguments={}, id="c1"), + ), + ChatMessage.from_user("current task"), + ] + generator, prompts = summarizer("summary") + SummarizationCompactor(generator, max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + # The summary cannot reproduce either attachment, so the transcript has to name them well enough to ask again. + assert "" in prompts[0] + assert "" in prompts[0] + + def test_custom_summary_instruction_replaces_the_default(self): + generator, prompts = summarizer("summary") + SummarizationCompactor(generator, summary_instruction="Only list file paths.", max_summary_tokens=1).compact( + messages=two_turns_and_a_task(), target_tokens=SMALLEST, token_counter=COUNTER + ) + assert "Only list file paths." in prompts[0] + assert "You are compacting part of a conversation" not in prompts[0] + + def test_keeps_partial_progress_when_a_summary_fails(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("old question " * 30), + ChatMessage.from_assistant("old answer " * 30), + ChatMessage.from_user("current task"), + ChatMessage.from_assistant("old step " * 30), + ChatMessage.from_assistant("new step"), + ] + generator, prompts = summarizer("history", RuntimeError("provider unavailable")) + + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + + assert compacted is not None + assert len(prompts) == 2 + # The history was summarized before the step summary failed, and that progress is kept. + assert sources(compacted) == ["historical_turns"] + assert compacted[-2:] == messages[-2:] + + def test_raises_when_a_summary_does_not_shrink_the_conversation(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("old"), + ChatMessage.from_assistant("answer"), + ChatMessage.from_user("current"), + ] + compactor = SummarizationCompactor( + MockChatGenerator("much longer summary " * 100), max_summary_tokens=1, raise_on_failure=True + ) + with pytest.raises(RuntimeError, match="did not reduce"): + compactor.compact(messages=messages, target_tokens=SMALLEST, token_counter=COUNTER) + + def test_returns_none_when_the_conversation_fits(self): + generator, prompts = summarizer("unused") + messages = [ChatMessage.from_system("rules"), ChatMessage.from_user("task")] + compacted = SummarizationCompactor(generator).compact( + messages=messages, target_tokens=10_000, token_counter=COUNTER + ) + assert compacted is None + assert prompts == [] + + def test_summary_is_a_marked_user_message(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("old " * 100), + ChatMessage.from_assistant("answer " * 100), + ChatMessage.from_user("task"), + ] + compacted = SummarizationCompactor(MockChatGenerator("summary"), max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + assert compacted is not None + assert compacted[1].is_from(role=ChatRole.USER) + assert compacted[1].meta[_COMPACTION_META_KEY] == { + "strategy": "summarization", + "summarized_messages": 2, + "source": "historical_turns", + } + + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"min_keep_steps": -1}, "`min_keep_steps` must be at least 0"), + ({"max_summary_tokens": 0}, "`max_summary_tokens` must be a positive"), + ], + ) + def test_rejects_invalid_settings(self, kwargs, match): + with pytest.raises(ValueError, match=match): + SummarizationCompactor(MockChatGenerator("summary"), **kwargs) + + def test_serde_round_trip(self): + compactor = SummarizationCompactor( + MockChatGenerator("summary"), + min_keep_steps=2, + max_summary_tokens=321, + summary_instruction="custom", + raise_on_failure=True, + ) + restored = SummarizationCompactor.from_dict(compactor.to_dict()) + assert isinstance(restored.chat_generator, MockChatGenerator) + assert restored.min_keep_steps == 2 + assert restored.max_summary_tokens == 321 + assert restored.summary_instruction == "custom" + assert restored.raise_on_failure is True + + +class TestSummarizationCompactorAsync: + @pytest.mark.asyncio + async def test_compact_async_matches_compact(self): + messages = two_turns_and_a_task() + generator, prompts = summarizer("async summary") + + compacted = await SummarizationCompactor(generator, max_summary_tokens=1).compact_async( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + + assert len(prompts) == 1 + assert compacted == SummarizationCompactor(MockChatGenerator("async summary"), max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) diff --git a/test/token_counters/test_utils.py b/test/token_counters/test_utils.py index b3101da360b..8fb06a896a2 100644 --- a/test/token_counters/test_utils.py +++ b/test/token_counters/test_utils.py @@ -75,6 +75,17 @@ def test_joins_messages_with_newlines(self): def test_empty_conversation(self): assert _rendered_conversation([]) == "" + def test_a_custom_placeholder_reaches_nested_tool_results(self): + messages = [ + ChatMessage.from_user(content_parts=["look:", IMAGE]), + ChatMessage.from_tool( + tool_result=[TextContent(text="screenshot: "), IMAGE], + origin=ToolCall(tool_name="browse", arguments={}, id="c1"), + ), + ] + rendered = _rendered_conversation(messages, placeholder=lambda content: "") + assert rendered == "[user] look:\n[user] \n[tool:browse] screenshot: " + @tool def search(query: Annotated[str, "the search query"]) -> str: