diff --git a/python/semantic_kernel/contents/history_reducer/chat_history_reducer_utils.py b/python/semantic_kernel/contents/history_reducer/chat_history_reducer_utils.py index a403cb21570a..cfb25fccc258 100644 --- a/python/semantic_kernel/contents/history_reducer/chat_history_reducer_utils.py +++ b/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 diff --git a/python/tests/unit/contents/test_chat_history_reducer_utils.py b/python/tests/unit/contents/test_chat_history_reducer_utils.py index d46340915b8a..25c2f2fbe429 100644 --- a/python/tests/unit/contents/test_chat_history_reducer_utils.py +++ b/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