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

python/semantic_kernel/contents/history_reducer/chat_history_reducer_utils.py

@@ -169,14 +169,14 @@ 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] = []
     i = 0
@@ -196,31 +196,22 @@ def extract_range(
 
         # 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]
+            paired_indices = sorted(p for p in pair_map[idx] if start <= p < end)
             # 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
+            if paired_indices:
+                # Check if any member of the pair or itself fails filter_func
+                if filter_func and (filter_func(msg) or any(filter_func(history[p]) for p in paired_indices)):
+                    # skip all
                     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)
+                    # Also skip the paired indices if they're in our current slice
+                    for paired_idx in paired_indices:
+                        if paired_idx > idx and paired_idx in sliced:
+                            # remove it from the slice so we don't process it again
+                            sliced.remove(paired_idx)
                     continue
-                # keep both
+                # keep all: each paired message is appended when its own index is reached,
+                # so messages interleaved between a call and its results keep their order
                 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
                 i += 1
                 continue
             # If the paired_idx is outside [start, end), there's no conflict

python/tests/unit/contents/test_chat_history_reducer_utils.py

@@ -230,3 +230,53 @@ 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