fix: merge streaming tool deltas by logical index

src/openai/lib/streaming/_assistants.py

@@ -5,7 +5,8 @@
 from typing import TYPE_CHECKING, Any, Generic, TypeVar, Callable, Iterable, Iterator, cast
 from typing_extensions import Awaitable, AsyncIterable, AsyncIterator, assert_never
 
-from ..._utils import is_dict, is_list, consume_sync_iterator, consume_async_iterator
+from ._deltas import accumulate_delta
+from ..._utils import consume_sync_iterator, consume_async_iterator
 from ..._compat import model_dump
 from ..._httpx2 import timeout_exceptions
 from ..._models import construct_type
@@ -978,64 +979,3 @@ def accumulate_event(
                 )
 
     return current_message_snapshot, new_content
-
-
-def accumulate_delta(acc: dict[object, object], delta: dict[object, object]) -> dict[object, object]:
-    for key, delta_value in delta.items():
-        if key not in acc:
-            acc[key] = delta_value
-            continue
-
-        acc_value = acc[key]
-        if acc_value is None:
-            acc[key] = delta_value
-            continue
-
-        # the `index` property is used in arrays of objects so it should
-        # not be accumulated like other values e.g.
-        # [{'foo': 'bar', 'index': 0}]
-        #
-        # the same applies to `type` properties as they're used for
-        # discriminated unions
-        if key == "index" or key == "type":
-            acc[key] = delta_value
-            continue
-
-        if isinstance(acc_value, str) and isinstance(delta_value, str):
-            acc_value += delta_value
-        elif isinstance(acc_value, (int, float)) and isinstance(delta_value, (int, float)):
-            acc_value += delta_value
-        elif is_dict(acc_value) and is_dict(delta_value):
-            acc_value = accumulate_delta(acc_value, delta_value)
-        elif is_list(acc_value) and is_list(delta_value):
-            # for lists of non-dictionary items we'll only ever get new entries
-            # in the array, existing entries will never be changed
-            if all(isinstance(x, (str, int, float)) for x in acc_value):
-                acc_value.extend(delta_value)
-                continue
-
-            for delta_entry in delta_value:
-                if not is_dict(delta_entry):
-                    raise TypeError(f"Unexpected list delta entry is not a dictionary: {delta_entry}")
-
-                try:
-                    index = delta_entry["index"]
-                except KeyError as exc:
-                    raise RuntimeError(f"Expected list delta entry to have an `index` key; {delta_entry}") from exc
-
-                if not isinstance(index, int):
-                    raise TypeError(f"Unexpected, list delta entry `index` value is not an integer; {index}")
-
-                try:
-                    acc_entry = acc_value[index]
-                except IndexError:
-                    acc_value.insert(index, delta_entry)
-                else:
-                    if not is_dict(acc_entry):
-                        raise TypeError("not handled yet")
-
-                    acc_value[index] = accumulate_delta(acc_entry, delta_entry)
-
-        acc[key] = acc_value
-
-    return acc

src/openai/lib/streaming/_deltas.py

@@ -6,11 +6,19 @@
 def accumulate_delta(acc: dict[object, object], delta: dict[object, object]) -> dict[object, object]:
     for key, delta_value in delta.items():
         if key not in acc:
+            if is_list(delta_value) and _has_indexed_entries(delta_value):
+                acc[key] = _accumulate_list_delta([], delta_value)
+                continue
+
             acc[key] = delta_value
             continue
 
         acc_value = acc[key]
         if acc_value is None:
+            if is_list(delta_value) and _has_indexed_entries(delta_value):
+                acc[key] = _accumulate_list_delta([], delta_value)
+                continue
+
             acc[key] = delta_value
             continue
 
@@ -31,34 +39,72 @@ def accumulate_delta(acc: dict[object, object], delta: dict[object, object]) ->
         elif is_dict(acc_value) and is_dict(delta_value):
             acc_value = accumulate_delta(acc_value, delta_value)
         elif is_list(acc_value) and is_list(delta_value):
-            # for lists of non-dictionary items we'll only ever get new entries
-            # in the array, existing entries will never be changed
-            if all(isinstance(x, (str, int, float)) for x in acc_value):
-                acc_value.extend(delta_value)
-                continue
+            acc_value = _accumulate_list_delta(acc_value, delta_value)
 
-            for delta_entry in delta_value:
-                if not is_dict(delta_entry):
-                    raise TypeError(f"Unexpected list delta entry is not a dictionary: {delta_entry}")
+        acc[key] = acc_value
 
-                try:
-                    index = delta_entry["index"]
-                except KeyError as exc:
-                    raise RuntimeError(f"Expected list delta entry to have an `index` key; {delta_entry}") from exc
+    return acc
 
-                if not isinstance(index, int):
-                    raise TypeError(f"Unexpected, list delta entry `index` value is not an integer; {index}")
 
-                try:
-                    acc_entry = acc_value[index]
-                except IndexError:
-                    acc_value.insert(index, delta_entry)
-                else:
-                    if not is_dict(acc_entry):
-                        raise TypeError("not handled yet")
+def _has_indexed_entries(value: list[object]) -> bool:
+    return any(is_dict(entry) and "index" in entry for entry in value)
 
-                    acc_value[index] = accumulate_delta(acc_entry, delta_entry)
 
-        acc[key] = acc_value
+def _accumulate_list_delta(acc_value: list[object], delta_value: list[object]) -> list[object]:
+    # for lists of non-dictionary items we'll only ever get new entries
+    # in the array, existing entries will never be changed
+    if not _has_indexed_entries(delta_value) and all(isinstance(x, (str, int, float)) for x in acc_value):
+        acc_value.extend(delta_value)
+        return acc_value
 
-    return acc
+    for delta_entry in delta_value:
+        if not is_dict(delta_entry):
+            raise TypeError(f"Unexpected list delta entry is not a dictionary: {delta_entry}")
+
+        try:
+            index = delta_entry["index"]
+        except KeyError as exc:
+            raise RuntimeError(f"Expected list delta entry to have an `index` key; {delta_entry}") from exc
+
+        if not isinstance(index, int):
+            raise TypeError(f"Unexpected, list delta entry `index` value is not an integer; {index}")
+
+        acc_index = _find_entry_index(acc_value, index)
+        if acc_index is None:
+            acc_value.insert(_find_insert_position(acc_value, index), delta_entry)
+            continue
+
+        acc_entry = acc_value[acc_index]
+        if not is_dict(acc_entry):
+            raise TypeError("not handled yet")
+
+        acc_value[acc_index] = accumulate_delta(acc_entry, delta_entry)
+
+    return acc_value
+
+
+def _find_entry_index(entries: list[object], index: int) -> int | None:
+    for entry_index, entry in enumerate(entries):
+        if is_dict(entry) and entry.get("index") == index:
+            return entry_index
+
+    # Full Assistants snapshots omit the delta-only `index` field.
+    # Preserve their established positional merge behavior as a fallback.
+    if index < len(entries):
+        positional_entry = entries[index]
+        if is_dict(positional_entry) and "index" not in positional_entry:
+            return index
+
+    return None
+
+
+def _find_insert_position(entries: list[object], index: int) -> int:
+    for entry_index, entry in enumerate(entries):
+        if not is_dict(entry):
+            continue
+
+        entry_delta_index = entry.get("index")
+        if isinstance(entry_delta_index, int) and entry_delta_index > index:
+            return entry_index
+
+    return len(entries)

src/openai/lib/streaming/chat/_completions.py

@@ -742,9 +742,13 @@ def _convert_initial_chunk_into_snapshot(chunk: ChatCompletionChunk) -> ParsedCh
     choices = cast("list[object]", data["choices"])
 
     for choice in chunk.choices:
+        message = accumulate_delta(
+            {},
+            cast("dict[object, object]", choice.delta.to_dict()),
+        )
         choices[choice.index] = {
             **choice.model_dump(exclude_unset=True, exclude={"delta"}),
-            "message": choice.delta.to_dict(),
+            "message": message,
         }
 
     return cast(

tests/lib/test_streaming_deltas.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+from typing import cast
+
+from openai.types.chat import ChatCompletionChunk
+from openai.lib.streaming.chat import ChatCompletionStreamState
+from openai.lib.streaming._deltas import accumulate_delta as accumulate_chat_delta
+from openai.lib.streaming._assistants import accumulate_delta as accumulate_assistant_delta
+from openai.types.chat.chat_completion_chunk import (
+    Choice,
+    ChoiceDelta,
+    ChoiceDeltaToolCall,
+    ChoiceDeltaToolCallFunction,
+)
+
+
+def test_accumulate_delta_merges_duplicate_indexed_entries_on_initial_chunk() -> None:
+    acc: dict[object, object] = {"tool_calls": None}
+
+    accumulate_chat_delta(
+        acc,
+        {
+            "tool_calls": [
+                {
+                    "index": 0,
+                    "id": "call_abc",
+                    "function": {"name": "get_weather"},
+                    "type": "function",
+                },
+                {"index": 0, "function": {"arguments": '{"city"'}},
+            ]
+        },
+    )
+    accumulate_chat_delta(acc, {"tool_calls": [{"index": 0, "function": {"arguments": ': "London"}'}}]})
+
+    assert acc == {
+        "tool_calls": [
+            {
+                "index": 0,
+                "id": "call_abc",
+                "function": {"name": "get_weather", "arguments": '{"city": "London"}'},
+                "type": "function",
+            }
+        ]
+    }
+
+
+def test_chat_completion_state_merges_duplicate_indexed_entries_on_initial_chunk() -> None:
+    state = ChatCompletionStreamState()
+
+    state.handle_chunk(
+        ChatCompletionChunk(
+            id="chatcmpl_abc",
+            choices=[
+                Choice(
+                    delta=ChoiceDelta(
+                        role="assistant",
+                        tool_calls=[
+                            ChoiceDeltaToolCall(
+                                index=0,
+                                id="call_abc",
+                                function=ChoiceDeltaToolCallFunction(name="get_weather", arguments='{"city"'),
+                                type="function",
+                            ),
+                            ChoiceDeltaToolCall(
+                                index=0,
+                                function=ChoiceDeltaToolCallFunction(arguments=': "London"}'),
+                            ),
+                        ],
+                    ),
+                    finish_reason=None,
+                    index=0,
+                    logprobs=None,
+                )
+            ],
+            created=1,
+            model="gpt-test",
+            object="chat.completion.chunk",
+        )
+    )
+    state.handle_chunk(
+        ChatCompletionChunk(
+            id="chatcmpl_abc",
+            choices=[
+                Choice(delta=ChoiceDelta(), finish_reason="tool_calls", index=0, logprobs=None),
+            ],
+            created=1,
+            model="gpt-test",
+            object="chat.completion.chunk",
+        )
+    )
+
+    tool_calls = state.get_final_completion().choices[0].message.tool_calls
+    assert tool_calls is not None
+    assert len(tool_calls) == 1
+    assert tool_calls[0].id == "call_abc"
+    assert tool_calls[0].function.name == "get_weather"
+    assert tool_calls[0].function.arguments == '{"city": "London"}'
+
+
+def test_assistant_accumulate_delta_uses_logical_index_for_initial_chunk() -> None:
+    acc: dict[object, object] = {}
+
+    accumulate_assistant_delta(
+        acc,
+        {
+            "tool_calls": [
+                {"index": 0, "id": "call_abc", "function": {"name": "get_weather"}, "type": "function"},
+                {"index": 0, "function": {"arguments": '{"path"'}},
+                {"index": 1, "id": "call_def", "function": {"name": "list_files"}, "type": "function"},
+            ]
+        },
+    )
+    accumulate_assistant_delta(
+        acc,
+        {
+            "tool_calls": [
+                {"index": 1, "function": {"arguments": '{"limit": 10}'}},
+                {"index": 0, "function": {"arguments": ': "."}'}},
+            ]
+        },
+    )
+
+    assert acc == {
+        "tool_calls": [
+            {
+                "index": 0,
+                "id": "call_abc",
+                "function": {"name": "get_weather", "arguments": '{"path": "."}'},
+                "type": "function",
+            },
+            {
+                "index": 1,
+                "id": "call_def",
+                "function": {"name": "list_files", "arguments": '{"limit": 10}'},
+                "type": "function",
+            },
+        ]
+    }
+
+
+def test_assistant_accumulate_delta_merges_indexed_delta_into_full_snapshot() -> None:
+    acc: dict[object, object] = {
+        "tool_calls": [
+            {
+                "id": "call_abc",
+                "function": {"name": "get_weather", "arguments": ""},
+                "type": "function",
+            }
+        ]
+    }
+
+    accumulate_assistant_delta(
+        acc,
+        {"tool_calls": [{"index": 0, "function": {"arguments": '{"city": "London"}'}}]},
+    )
+
+    tool_calls = cast(list[object], acc["tool_calls"])
+    assert len(tool_calls) == 1
+    assert tool_calls[0] == {
+        "index": 0,
+        "id": "call_abc",
+        "function": {"name": "get_weather", "arguments": '{"city": "London"}'},
+        "type": "function",
+    }