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 84e79c29ecc1..2e0e82fad8ab 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 @@ -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 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..9fa6ce2ef4ce 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,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