Python: Keep parallel function call/result groups together in extract_range

python/semantic_kernel/contents/history_reducer/chat_history_reducer_utils.py

@@ -205,21 +205,26 @@ def extract_range(
     sliced = list(range(start, end))
 
     # If we need to preserve call->result pairs, gather them
-    pair_map = {}
+    pair_map: dict[int, set[int]] = {}
     if preserve_pairs:
         pairs = get_call_result_pairs(history)
         # store in a dict for quick membership checking
-        # call_idx -> result_idx, and also result_idx -> call_idx
+        # call_idx -> result_idxs (several for parallel calls), and also result_idx -> {call_idx}
         for cidx, ridx in pairs:
-            pair_map[cidx] = ridx
-            pair_map[ridx] = cidx
+            pair_map.setdefault(cidx, set()).add(ridx)
+            pair_map.setdefault(ridx, set()).add(cidx)
 
     extracted: list[ChatMessageContent] = []
+    dropped: set[int] = set()
     i = 0
     while i < len(sliced):
         idx = sliced[i]
         msg = history[idx]
 
+        if idx in dropped:
+            i += 1
+            continue
+
         # If filter_func excludes it, skip it
         if filter_func and filter_func(msg):
             i += 1
@@ -230,43 +235,26 @@ def extract_range(
             i += 1
             continue
 
-        # If preserve_pairs is on, and there's a paired index, skip or include them both
-        if preserve_pairs and idx in pair_map:
-            paired_idx = pair_map[idx]
-            # If the pair is within [start, end), we must keep or skip them together
-            if start <= paired_idx < end:
-                # Check if the pair or itself fails filter_func
-                if filter_func and (filter_func(history[paired_idx]) or filter_func(msg)):
-                    # skip both
-                    i += 1
-                    # Also skip the paired index if it's in our current slice
-                    if paired_idx in sliced:
-                        # remove it from the slice so we don't process it again
-                        sliced.remove(paired_idx)
-                    continue
-                # keep both
-                extracted.append(msg)
-                if paired_idx > idx:
-                    # We'll skip the pair in the normal iteration by removing from slice
-                    # but add it to extracted right now
-                    extracted.append(history[paired_idx])
-                    if paired_idx in sliced:
-                        sliced.remove(paired_idx)
-                else:
-                    # if paired_idx < idx, it might appear later, so skip for now
-                    # but we may have already processed it if i was the 2nd item
-                    # either way, do not add duplicates
-                    pass
+        # If preserve_pairs is on, keep or skip the whole connected call/result group together
+        if preserve_pairs and idx in pair_map and filter_func:
+            # A message can link several calls (one message holding multiple results) and vice versa,
+            # so walk the connected group within [start, end) rather than only direct neighbours.
+            group = {idx}
+            stack = [idx]
+            while stack:
+                for p in pair_map.get(stack.pop(), ()):
+                    if start <= p < end and p not in group:
+                        group.add(p)
+                        stack.append(p)
+            if any(filter_func(history[p]) for p in group):
+                # skip all members; later ones are dropped when their index is reached
+                dropped.update(group)
                 i += 1
                 continue
-            # If the paired_idx is outside [start, end), there's no conflict
-            # so we can just do normal logic
-            extracted.append(msg)
-            i += 1
-        else:
-            # keep it if filter_func not triggered
-            extracted.append(msg)
-            i += 1
+        # each kept message is appended when its own index is reached, so messages
+        # interleaved between a call and its results keep their original order
+        extracted.append(msg)
+        i += 1
 
     return extracted
 

python/tests/unit/contents/test_chat_history_reducer_utils.py

@@ -230,3 +230,82 @@ def test_locate_safe_reduction_index_tool_role_without_function_result_content()
             f"Tool result at index 2 was kept but tool call at index 1 was dropped. "
             f"Kept indices: {kept_indices}, reduction index: {idx}"
         )
+
+
+def test_extract_range_preserve_pairs_parallel_calls_filtered_result():
+    """
+    When one assistant message holds several parallel function calls and filter_func drops only
+    one of their results, the call and its remaining results must not be split: either all are
+    kept or all are skipped, so no result is left without its call.
+    """
+    user = ChatMessageContent(role=AuthorRole.USER, content="weather?")
+    call = ChatMessageContent(
+        role=AuthorRole.ASSISTANT,
+        items=[
+            FunctionCallContent(id="call_paris", function_name="get_weather", arguments={"city": "Paris"}),
+            FunctionCallContent(id="call_tokyo", function_name="get_weather", arguments={"city": "Tokyo"}),
+        ],
+    )
+    result_paris = ChatMessageContent(
+        role=AuthorRole.TOOL, items=[FunctionResultContent(id="call_paris", function_name="get_weather", result="15C")]
+    )
+    result_tokyo = ChatMessageContent(
+        role=AuthorRole.TOOL, items=[FunctionResultContent(id="call_tokyo", function_name="get_weather", result="22C")]
+    )
+    done = ChatMessageContent(role=AuthorRole.ASSISTANT, content="done")
+    history = [user, call, result_paris, result_tokyo, done]
+
+    assert get_call_result_pairs(history) == [(1, 2), (1, 3)]
+
+    extracted = extract_range(
+        history, start=0, end=len(history), filter_func=lambda m: m is result_tokyo, preserve_pairs=True
+    )
+    assert extracted == [user, done]
+
+    extracted = extract_range(history, start=0, end=len(history), preserve_pairs=True)
+    assert extracted == history
+
+    # Filtering the call itself skips all of its results, without skipping unrelated messages
+    extracted = extract_range(history, start=0, end=len(history), filter_func=lambda m: m is call, preserve_pairs=True)
+    assert extracted == [user, done]
+
+    # With the call outside the range, a remaining result is not tied to its filtered sibling
+    extracted = extract_range(
+        history, start=2, end=len(history), filter_func=lambda m: m is result_paris, preserve_pairs=True
+    )
+    assert extracted == [result_tokyo, done]
+
+    # A message interleaved between the call and its results keeps its position
+    note = ChatMessageContent(role=AuthorRole.USER, content="also Tokyo")
+    history = [call, note, result_paris, result_tokyo]
+    extracted = extract_range(history, start=0, end=len(history), preserve_pairs=True)
+    assert extracted == history
+
+
+def test_extract_range_preserve_pairs_shared_result_message_filtered_call():
+    """
+    One tool message holding the results of two calls links both calls together, so filtering
+    out either call must drop the shared result and the other call too.
+    """
+    call_a = ChatMessageContent(
+        role=AuthorRole.ASSISTANT, items=[FunctionCallContent(id="a", function_name="f", arguments={})]
+    )
+    call_b = ChatMessageContent(
+        role=AuthorRole.ASSISTANT, items=[FunctionCallContent(id="b", function_name="f", arguments={})]
+    )
+    results = ChatMessageContent(
+        role=AuthorRole.TOOL,
+        items=[
+            FunctionResultContent(id="a", function_name="f", result="1"),
+            FunctionResultContent(id="b", function_name="f", result="2"),
+        ],
+    )
+    history = [call_a, call_b, results]
+
+    assert get_call_result_pairs(history) == [(0, 2), (1, 2)]
+
+    extracted = extract_range(history, start=0, end=3, filter_func=lambda m: m is call_a, preserve_pairs=True)
+    assert extracted == []
+
+    extracted = extract_range(history, start=0, end=3, preserve_pairs=True)
+    assert extracted == history