From dcbc949405d556057e924580bc88d483378d8963 Mon Sep 17 00:00:00 2001 From: King Star Date: Wed, 22 Jul 2026 03:12:22 +0800 Subject: [PATCH] Python: Keep non-adjacent function call pairs atomic --- .../core/agent_framework/_compaction.py | 145 ++++++++++++- .../core/tests/core/test_compaction.py | 197 ++++++++++++++++++ 2 files changed, 333 insertions(+), 9 deletions(-) diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index 59abb10a468..485d2c93ce9 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -101,6 +101,18 @@ def _is_reasoning_only_assistant(message: Message) -> bool: return all(content.type == "text_reasoning" for content in message.contents) +def _function_call_ids(message: Message) -> set[str]: + if message.role != "assistant": + return set() + return {content.call_id for content in message.contents if content.type == "function_call" and content.call_id} + + +def _function_result_ids(message: Message) -> set[str]: + if message.role != "tool": + return set() + return {content.call_id for content in message.contents if content.type == "function_result" and content.call_id} + + def _ensure_message_ids( messages: list[Message], *, id_offset: int = 0, reserved_ids: Iterable[str] | None = None ) -> None: @@ -126,6 +138,68 @@ def _group_id_for(message: Message, group_index: int) -> str: return f"group_index_{group_index}" +def _link_function_call_result_spans(messages: Sequence[Message], spans: list[dict[str, Any]]) -> None: + """Link non-adjacent function results to their unique declaration group.""" + if len(spans) < 2: + return + + declaration_spans: dict[str, set[int]] = {} + result_ids_by_span: list[set[str]] = [] + for span_index, span in enumerate(spans): + start_index = int(span["start_index"]) + end_index = int(span["end_index"]) + declared_ids: set[str] = set() + result_ids: set[str] = set() + for message in messages[start_index : end_index + 1]: + declared_ids.update(_function_call_ids(message)) + result_ids.update(_function_result_ids(message)) + for call_id in declared_ids: + declaration_spans.setdefault(call_id, set()).add(span_index) + result_ids_by_span.append(result_ids) + if not declaration_spans or not any(result_ids_by_span): + return + + parents = list(range(len(spans))) + + def find(span_index: int) -> int: + while parents[span_index] != span_index: + parents[span_index] = parents[parents[span_index]] + span_index = parents[span_index] + return span_index + + def union(left: int, right: int) -> bool: + left_root = find(left) + right_root = find(right) + if left_root == right_root: + return False + earlier_root = min(left_root, right_root) + later_root = max(left_root, right_root) + parents[later_root] = earlier_root + return True + + linked = False + for result_span_index, result_ids in enumerate(result_ids_by_span): + for call_id in result_ids: + matches = declaration_spans.get(call_id) + if matches is None or len(matches) != 1: + continue + declaration_span_index = next(iter(matches)) + if declaration_span_index < result_span_index and union(result_span_index, declaration_span_index): + linked = True + if not linked: + return + + has_reasoning_by_root: dict[int, bool] = {} + for span_index, span in enumerate(spans): + root = find(span_index) + has_reasoning_by_root[root] = has_reasoning_by_root.get(root, False) or bool(span["has_reasoning"]) + + for span_index, span in enumerate(spans): + root = find(span_index) + span["group_id"] = spans[root]["group_id"] + span["has_reasoning"] = has_reasoning_by_root[root] + + def group_messages( messages: list[Message], *, id_offset: int = 0, reserved_ids: Iterable[str] | None = None ) -> list[dict[str, Any]]: @@ -145,6 +219,7 @@ def group_messages( Returns: Ordered list of lightweight span dicts with keys: ``group_id``, ``kind``, ``start_index``, ``end_index``, ``has_reasoning``. + Non-contiguous function-call declaration and result spans share a group id. """ _ensure_message_ids(messages, id_offset=id_offset, reserved_ids=reserved_ids) spans: list[dict[str, Any]] = [] @@ -249,6 +324,7 @@ def group_messages( i += 1 group_index += 1 + _link_function_call_result_spans(messages, spans) return spans @@ -434,6 +510,38 @@ def _reannotation_start(messages: Sequence[Message], index: int) -> int: return previous_index +def _function_pair_reannotation_start(messages: Sequence[Message], start_index: int) -> int: + result_ids: set[str] = set() + for message in messages[start_index:]: + result_ids.update(_function_result_ids(message)) + if not result_ids: + return start_index + + declaration_indices: dict[str, set[int]] = {} + for index, message in enumerate(messages): + for call_id in _function_call_ids(message): + declaration_indices.setdefault(call_id, set()).add(index) + + matching_indices: list[int] = [] + for call_id in result_ids: + indices = declaration_indices.get(call_id) + if indices is None or len(indices) != 1: + continue + declaration_index = next(iter(indices)) + if declaration_index < start_index: + matching_indices.append(declaration_index) + if not matching_indices: + return start_index + + earliest_index = min(matching_indices) + declaration_group_id = _group_id(messages[earliest_index]) + if declaration_group_id is None: + return earliest_index + while earliest_index > 0 and _group_id(messages[earliest_index - 1]) == declaration_group_id: + earliest_index -= 1 + return earliest_index + + def annotate_message_groups( messages: list[Message], *, @@ -444,7 +552,8 @@ def annotate_message_groups( """Annotate message groups while reusing existing annotations when possible. By default, the function re-annotates only the suffix that contains new - messages and keeps previously annotated prefixes untouched. When a + messages and keeps previously annotated prefixes untouched. A newly added + function result expands that suffix back to its unique declaration. When a ``tokenizer`` is provided, token-count annotations are also populated incrementally. """ @@ -466,18 +575,29 @@ def annotate_message_groups( start_index = min(candidate_starts) start_index = _reannotation_start(messages, start_index) + start_index = _function_pair_reannotation_start(messages, start_index) - # Continue group indices from the preserved prefix when only re-annotating a suffix. - group_index_offset = 0 - if start_index > 0: - previous_group_index = _group_index(messages[start_index - 1]) - if previous_group_index is not None: - group_index_offset = previous_group_index + 1 + # Linked groups can be non-contiguous, so the last prefix message does not + # necessarily carry the highest group index. + prefix_group_indices = [ + group_index for message in messages[:start_index] if (group_index := _group_index(message)) is not None + ] + group_index_offset = max(prefix_group_indices, default=-1) + 1 reserved_ids = {message.message_id for message in messages[:start_index] if message.message_id} spans = group_messages(messages[start_index:], id_offset=start_index, reserved_ids=reserved_ids) - for span_index, span in enumerate(spans): + span_counts_by_group_id: dict[str, int] = {} + for span in spans: + group_id = str(span["group_id"]) + span_counts_by_group_id[group_id] = span_counts_by_group_id.get(group_id, 0) + 1 + linked_group_ids = {group_id for group_id, count in span_counts_by_group_id.items() if count > 1} + + group_indices: dict[str, int] = {} + grouped_messages: dict[str, list[Message]] = {} + for span in spans: group_id = str(span["group_id"]) + if group_id not in group_indices: + group_indices[group_id] = group_index_offset + len(group_indices) kind = _coerce_group_kind(span["kind"]) if kind is None: raise ValueError(f"Unexpected group kind in span: {span['kind']}") @@ -490,12 +610,19 @@ def annotate_message_groups( message, group_id=group_id, kind=kind, - index=group_index_offset + span_index, + index=group_indices[group_id], has_reasoning=has_reasoning, ) message.additional_properties.setdefault(EXCLUDED_KEY, False) + if group_id in linked_group_ids: + grouped_messages.setdefault(group_id, []).append(message) if tokenizer is not None and _token_count(message) is None: _write_token_count(message, tokenizer.count_tokens(_serialize_message(message))) + + for group in grouped_messages.values(): + if any(not message.additional_properties.get(EXCLUDED_KEY, False) for message in group): + for message in group: + message.additional_properties[EXCLUDED_KEY] = False return _ordered_group_ids_from_annotations(messages) diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index 4507d39946f..b8a4b9e96fd 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -12,6 +12,7 @@ GROUP_ANNOTATION_KEY, GROUP_HAS_REASONING_KEY, GROUP_ID_KEY, + GROUP_INDEX_KEY, GROUP_KIND_KEY, GROUP_TOKEN_COUNT_KEY, SUMMARIZED_BY_SUMMARY_ID_KEY, @@ -113,6 +114,14 @@ def _group_kind(message: Message) -> str | None: return value if isinstance(value, str) else None +def _group_index(message: Message) -> int | None: + annotation = message.additional_properties.get(GROUP_ANNOTATION_KEY) + if not isinstance(annotation, dict): + return None + value = annotation.get(GROUP_INDEX_KEY) + return value if isinstance(value, int) else None + + def _group_has_reasoning(message: Message) -> bool | None: annotation = message.additional_properties.get(GROUP_ANNOTATION_KEY) if not isinstance(annotation, dict): @@ -186,6 +195,107 @@ def test_group_annotations_handle_same_message_reasoning_and_function_calls() -> assert _group_has_reasoning(messages[1]) is True +def test_group_annotations_pair_nonadjacent_function_result_by_call_id() -> None: + messages = [ + _assistant_reasoning_and_function_calls("c1"), + Message(role="assistant", contents=["approval completed"]), + _tool_result("c1", "ok"), + ] + + annotate_message_groups(messages) + + call_group = _group_id(messages[0]) + assert call_group is not None + assert _group_id(messages[2]) == call_group + assert _group_index(messages[2]) == _group_index(messages[0]) + assert _group_has_reasoning(messages[2]) is True + assert _group_id(messages[1]) != call_group + + +def test_group_annotations_pair_multiple_nonadjacent_results_with_declaration() -> None: + messages = [ + _assistant_reasoning_and_function_calls("c1", "c2"), + Message(role="assistant", contents=["first approval"]), + _tool_result("c1", "ok1"), + Message(role="assistant", contents=["second approval"]), + _tool_result("c2", "ok2"), + ] + + annotate_message_groups(messages) + + call_group = _group_id(messages[0]) + assert call_group is not None + assert _group_id(messages[2]) == call_group + assert _group_id(messages[4]) == call_group + assert _group_index(messages[2]) == _group_index(messages[0]) + assert _group_index(messages[4]) == _group_index(messages[0]) + + +def test_group_annotations_merge_declaration_groups_for_combined_result_message() -> None: + messages = [ + _assistant_function_call("c1"), + Message(role="assistant", contents=["between calls"]), + _assistant_function_call("c2"), + Message(role="assistant", contents=["approval completed"]), + Message( + role="tool", + contents=[ + Content.from_function_result(call_id="c1", result="ok1"), + Content.from_function_result(call_id="c2", result="ok2"), + ], + ), + ] + + annotate_message_groups(messages) + + call_group = _group_id(messages[0]) + assert call_group is not None + assert _group_id(messages[2]) == call_group + assert _group_id(messages[4]) == call_group + assert _group_index(messages[2]) == _group_index(messages[0]) + assert _group_index(messages[4]) == _group_index(messages[0]) + + +def test_group_annotations_leave_unmatched_result_separate_from_pending_call() -> None: + messages = [ + _assistant_function_call("pending"), + Message(role="assistant", contents=["waiting"]), + _tool_result("unknown", "result"), + ] + + annotate_message_groups(messages) + + assert _group_id(messages[0]) != _group_id(messages[2]) + + +def test_group_annotations_do_not_pair_result_before_declaration() -> None: + messages = [ + _tool_result("late", "result"), + Message(role="assistant", contents=["between"]), + _assistant_function_call("late"), + ] + + annotate_message_groups(messages) + + assert _group_id(messages[0]) != _group_id(messages[2]) + + +def test_group_annotations_do_not_pair_ambiguous_duplicate_call_ids() -> None: + messages = [ + _assistant_function_call("duplicate"), + Message(role="assistant", contents=["between declarations"]), + _assistant_function_call("duplicate"), + Message(role="assistant", contents=["before result"]), + _tool_result("duplicate", "result"), + ] + + annotate_message_groups(messages) + + result_group = _group_id(messages[4]) + assert result_group != _group_id(messages[0]) + assert result_group != _group_id(messages[2]) + + async def test_sliding_window_keeps_reasoning_and_mcp_call_atomic() -> None: messages = [ Message(role="system", contents=["system"]), @@ -258,6 +368,64 @@ def test_extend_compaction_messages_preserves_existing_annotations_and_tokens() assert isinstance(_token_count(messages[1]), int) +def test_extend_compaction_messages_pairs_nonadjacent_result_incrementally() -> None: + tokenizer = CharacterEstimatorTokenizer() + messages = [ + _assistant_function_call("c4"), + Message(role="assistant", contents=["approval completed"]), + ] + annotate_message_groups(messages, tokenizer=tokenizer) + call_group = _group_id(messages[0]) + intervening_group = _group_id(messages[1]) + + extend_compaction_messages(messages, [_tool_result("c4", "ok")], tokenizer=tokenizer) + + assert _group_id(messages[0]) == call_group + assert _group_id(messages[1]) == intervening_group + assert _group_id(messages[2]) == call_group + assert _group_index(messages[2]) == _group_index(messages[0]) + assert isinstance(_token_count(messages[2]), int) + + append_compaction_message( + messages, + Message(role="assistant", contents=["final answer"]), + tokenizer=tokenizer, + ) + + assert _group_id(messages[3]) not in {call_group, intervening_group} + assert _group_index(messages[3]) == 2 + + +def test_extend_compaction_messages_reincludes_excluded_declaration_for_new_result() -> None: + messages = [ + _assistant_function_call("c5"), + Message(role="assistant", contents=["approval pending"]), + ] + annotate_message_groups(messages) + messages[0].additional_properties[EXCLUDED_KEY] = True + + extend_compaction_messages(messages, [_tool_result("c5", "ok")]) + + assert messages[0].additional_properties[EXCLUDED_KEY] is False + assert messages[2].additional_properties[EXCLUDED_KEY] is False + assert _group_id(messages[2]) == _group_id(messages[0]) + + +def test_extend_compaction_messages_preserves_adjacent_duplicate_call_pair() -> None: + messages = [ + _assistant_function_call("duplicate"), + Message(role="assistant", contents=["between declarations"]), + _assistant_function_call("duplicate"), + ] + annotate_message_groups(messages) + + extend_compaction_messages(messages, [_tool_result("duplicate", "result")]) + + result_group = _group_id(messages[3]) + assert result_group != _group_id(messages[0]) + assert result_group == _group_id(messages[2]) + + def test_append_compaction_message_annotates_new_message() -> None: messages = [Message(role="user", contents=["hello"])] annotate_message_groups(messages) @@ -383,6 +551,35 @@ async def test_truncation_strategy_keeps_latest_group_when_it_exceeds_target() - assert included_messages(messages) == messages +async def test_truncation_strategy_keeps_nonadjacent_tool_pair_atomic() -> None: + tokenizer = CharacterEstimatorTokenizer() + messages = [ + Message(role="user", contents=["original request " + "x" * 1600]), + _assistant_function_call("call-1"), + Message(role="assistant", contents=["intervening approval traffic " + "y" * 1600]), + _tool_result("call-1", "result"), + Message(role="user", contents=["follow up " + "z" * 400]), + ] + strategy = TruncationStrategy( + max_n=600, + compact_to=520, + tokenizer=tokenizer, + ) + annotate_message_groups(messages, tokenizer=tokenizer) + + changed = await strategy(messages) + + assert changed is True + projected = included_messages(messages) + declared = { + content.call_id for message in projected for content in message.contents if content.type == "function_call" + } + results = { + content.call_id for message in projected for content in message.contents if content.type == "function_result" + } + assert declared == results + + def test_truncation_strategy_validates_token_targets() -> None: try: TruncationStrategy(max_n=3, compact_to=4)