diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index ce47f5fc8e..9c165ba6d4 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -1973,11 +1973,80 @@ def _coalesce_text_content(contents: list[Content], type_str: Literal["text", "t contents.extend(coalesced_contents) +def _coalesce_code_interpreter_tool_calls(contents: list[Content]) -> None: + """Coalesce code_interpreter_tool_call items with the same call_id, keeping the most complete chunk. + + The winning chunk stays at its original position so that relative ordering + with non-CI items is preserved. + """ + best_idx: dict[str, int] = {} + drop_indices: set[int] = set() + for i, content in enumerate(contents): + if content.type != "code_interpreter_tool_call" or not content.call_id: + continue + cid = content.call_id + if cid in best_idx: + prev = best_idx[cid] + if _code_interpreter_chunk_is_more_complete(content, contents[prev]): + drop_indices.add(prev) + best_idx[cid] = i + else: + drop_indices.add(i) + else: + best_idx[cid] = i + if not drop_indices: + return + for idx in sorted(drop_indices, reverse=True): + contents.pop(idx) + + +def _get_ci_chunk_content_length(content: Content) -> int: + """Return the total text length across all inputs of a code_interpreter_tool_call.""" + if not content.inputs: + return 0 + total = 0 + for inp in content.inputs: + if inp.type == "text" and inp.text: + total += len(inp.text) + return total + + +def _code_interpreter_chunk_is_more_complete(a: Content, b: Content) -> bool: + """Return True if 'a' is more complete than 'b'. + + Comparison order: + 1. A chunk with a valid int sequence_number beats one without. + 2. Higher sequence_number wins when both have one. + 3. Fallback: longer total input text wins. + """ + seq_a = _try_parse_seq(a) + seq_b = _try_parse_seq(b) + if seq_a is not None and seq_b is None: + return True + if seq_b is not None and seq_a is None: + return False + if seq_a is not None and seq_b is not None: + return seq_a > seq_b + return _get_ci_chunk_content_length(a) > _get_ci_chunk_content_length(b) + + +def _try_parse_seq(content: Content) -> int | None: + """Extract and validate sequence_number from additional_properties.""" + val = content.additional_properties.get("sequence_number") + if val is None: + return None + try: + return int(val) + except (ValueError, TypeError): + return None + + def _finalize_response(response: ChatResponse | AgentResponse) -> None: """Finalizes the response by performing any necessary post-processing.""" for msg in response.messages: _coalesce_text_content(msg.contents, "text") _coalesce_text_content(msg.contents, "text_reasoning") + _coalesce_code_interpreter_tool_calls(msg.contents) # region ContinuationToken diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 62e88fab02..b3773c401c 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -1740,6 +1740,129 @@ def test_coalesce_text_reasoning_with_different_ids(): assert contents[1].text == "Thinking B1 B2" +def test_coalesce_code_interpreter_tool_calls_keeps_most_complete(): + """Test that code_interpreter_tool_call chunks with same call_id coalesce to the most complete one.""" + from agent_framework._types import _coalesce_code_interpreter_tool_calls + + c1 = Content.from_code_interpreter_tool_call(call_id="ci_1", inputs=[Content.from_text("import ")]) + c1.additional_properties["sequence_number"] = 1 + c2 = Content.from_code_interpreter_tool_call(call_id="ci_1", inputs=[Content.from_text("import pandas")]) + c2.additional_properties["sequence_number"] = 2 + + contents = [c1, c2] + _coalesce_code_interpreter_tool_calls(contents) + assert len(contents) == 1 + assert contents[0].inputs[0].text == "import pandas" + + +def test_coalesce_code_interpreter_tool_calls_groups_by_call_id(): + """Test that multiple distinct call_ids each keep their own winning chunk.""" + from agent_framework._types import _coalesce_code_interpreter_tool_calls + + c1 = Content.from_code_interpreter_tool_call(call_id="ci_a", inputs=[Content.from_text("a1")]) + c1.additional_properties["sequence_number"] = 1 + c2 = Content.from_code_interpreter_tool_call(call_id="ci_a", inputs=[Content.from_text("a2")]) + c2.additional_properties["sequence_number"] = 2 + c3 = Content.from_code_interpreter_tool_call(call_id="ci_b", inputs=[Content.from_text("b1")]) + c3.additional_properties["sequence_number"] = 1 + + contents = [c1, c2, c3] + _coalesce_code_interpreter_tool_calls(contents) + assert len(contents) == 2 + assert contents[0].inputs[0].text == "a2" + assert contents[1].inputs[0].text == "b1" + + +def test_coalesce_code_interpreter_tool_calls_preserves_non_ci_items(): + """Test that non-CI items are preserved during coalescing.""" + from agent_framework._types import _coalesce_code_interpreter_tool_calls + + text_before = Content.from_text("before") + ci1 = Content.from_code_interpreter_tool_call(call_id="ci_x", inputs=[Content.from_text("short")]) + ci1.additional_properties["sequence_number"] = 1 + ci2 = Content.from_code_interpreter_tool_call(call_id="ci_x", inputs=[Content.from_text("longer_code")]) + ci2.additional_properties["sequence_number"] = 2 + text_after = Content.from_text("after") + + contents = [text_before, ci1, ci2, text_after] + _coalesce_code_interpreter_tool_calls(contents) + assert len(contents) == 3 + assert contents[0].text == "before" + assert contents[2].text == "after" + + +def test_coalesce_code_interpreter_tool_calls_no_sequence_number(): + """Test fallback to longest text when sequence_number is absent.""" + from agent_framework._types import _coalesce_code_interpreter_tool_calls + + c1 = Content.from_code_interpreter_tool_call(call_id="ci_y", inputs=[Content.from_text("short")]) + c2 = Content.from_code_interpreter_tool_call(call_id="ci_y", inputs=[Content.from_text("longer_script")]) + + contents = [c1, c2] + _coalesce_code_interpreter_tool_calls(contents) + assert len(contents) == 1 + assert contents[0].inputs[0].text == "longer_script" + + +def test_coalesce_code_interpreter_tool_calls_single_call_is_noop(): + """Test that a single CI call is unchanged.""" + from agent_framework._types import _coalesce_code_interpreter_tool_calls + + c1 = Content.from_code_interpreter_tool_call(call_id="ci_z", inputs=[Content.from_text("print(1)")]) + c1.additional_properties["sequence_number"] = 1 + + contents = [c1] + _coalesce_code_interpreter_tool_calls(contents) + assert len(contents) == 1 + assert contents[0].inputs[0].text == "print(1)" + + +def test_coalesce_code_interpreter_tool_calls_non_contiguous(): + """Test that non-contiguous CI chunks with the same call_id are coalesced and ordering preserved.""" + from agent_framework._types import _coalesce_code_interpreter_tool_calls + + ci1 = Content.from_code_interpreter_tool_call(call_id="ci_x", inputs=[Content.from_text("short")]) + ci1.additional_properties["sequence_number"] = 1 + other = Content.from_text("interleaved") + ci2 = Content.from_code_interpreter_tool_call(call_id="ci_x", inputs=[Content.from_text("longer_code")]) + ci2.additional_properties["sequence_number"] = 2 + + contents = [ci1, other, ci2] + _coalesce_code_interpreter_tool_calls(contents) + assert len(contents) == 2 + assert contents[0].text == "interleaved" + assert contents[1].inputs[0].text == "longer_code" + + +def test_coalesce_code_interpreter_tool_calls_mixed_seq_presence(): + """Test that a chunk with sequence_number beats one without, regardless of text length.""" + from agent_framework._types import _coalesce_code_interpreter_tool_calls + + no_seq = Content.from_code_interpreter_tool_call(call_id="ci_y", inputs=[Content.from_text("longer_text_without_seq")]) + has_seq = Content.from_code_interpreter_tool_call(call_id="ci_y", inputs=[Content.from_text("short")]) + has_seq.additional_properties["sequence_number"] = 5 + + contents = [no_seq, has_seq] + _coalesce_code_interpreter_tool_calls(contents) + assert len(contents) == 1 + assert contents[0].inputs[0].text == "short" + + +def test_coalesce_code_interpreter_tool_calls_string_sequence_number(): + """Test that string-typed sequence_number values are handled correctly.""" + from agent_framework._types import _coalesce_code_interpreter_tool_calls + + c1 = Content.from_code_interpreter_tool_call(call_id="ci_z", inputs=[Content.from_text("first")]) + c1.additional_properties["sequence_number"] = "1" + c2 = Content.from_code_interpreter_tool_call(call_id="ci_z", inputs=[Content.from_text("second")]) + c2.additional_properties["sequence_number"] = "2" + + contents = [c1, c2] + _coalesce_code_interpreter_tool_calls(contents) + assert len(contents) == 1 + assert contents[0].inputs[0].text == "second" + + def test_comprehensive_to_dict_exclude_options(): """Test to_dict methods with various exclude options for better coverage."""