Fix streaming delta merge for duplicate tool call indexes

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 as 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

@@ -3,16 +3,22 @@
 from ..._utils import is_dict, is_list
 
 
+def _has_indexed_entries(value: object) -> bool:
+    return is_list(value) and any(is_dict(entry) and "index" in entry for entry in value)
+
+
 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
+        if key not in acc or acc[key] is None:
+            if _has_indexed_entries(delta_value):
+                # Even the first indexed list must be merged: one chunk can
+                # contain several fragments for the same index.
+                acc[key] = []
+            else:
+                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.
@@ -33,7 +39,11 @@ def accumulate_delta(acc: dict[object, object], delta: dict[object, object]) ->
         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):
+            # An empty accumulator cannot tell us whether this is an indexed
+            # list. Inspect the delta before taking the append-only path.
+            if all(isinstance(x, (str, int, float)) for x in acc_value) and (
+                acc_value or not _has_indexed_entries(delta_value)
+            ):
                 acc_value.extend(delta_value)
                 continue
 

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

@@ -415,7 +415,7 @@ def _accumulate_chunk(self, chunk: ChatCompletionChunk) -> ParsedChatCompletionS
                         type_=ParsedChoiceSnapshot,
                         value={
                             **choice.model_dump(exclude_unset=True, exclude={"delta"}),
-                            "message": choice.delta.to_dict(),
+                            "message": accumulate_delta({}, cast("dict[object, object]", choice.delta.to_dict())),
                         },
                     ),
                 )
@@ -745,7 +745,7 @@ def _convert_initial_chunk_into_snapshot(chunk: ChatCompletionChunk) -> ParsedCh
     for choice in chunk.choices:
         choices[choice.index] = {
             **choice.model_dump(exclude_unset=True, exclude={"delta"}),
-            "message": choice.delta.to_dict(),
+            "message": accumulate_delta({}, cast("dict[object, object]", choice.delta.to_dict())),
         }
 
     return cast(

tests/lib/chat/test_tool_call_deltas.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+import httpx2
+import pytest
+from pydantic import BaseModel
+
+import openai
+from openai import OpenAI, AsyncOpenAI
+from tests.respx2 import MockRouter
+from openai.types.chat import ParsedChatCompletion
+from openai.lib.streaming.chat import ChatCompletionStreamEvent
+
+from ...conftest import base_url
+
+
+class ListFilesArgs(BaseModel):
+    path: str
+
+
+@pytest.fixture(params=["initial", "after_role", "empty", "new_choice", "existing_tool", "parallel"])
+def tool_call_chunks(request: pytest.FixtureRequest) -> list[dict[str, Any]]:
+    scenario = request.param
+    chunks: list[dict[str, Any]] = []
+    choice_index = 1 if scenario == "new_choice" else 0
+    tool_index = 1 if scenario in ("existing_tool", "parallel") else 0
+
+    def chunk(delta: dict[str, Any], index: int = choice_index, finish: str | None = None) -> dict[str, Any]:
+        return {
+            "id": "chatcmpl-test",
+            "object": "chat.completion.chunk",
+            "created": 0,
+            "model": "gpt-4o-mini",
+            "choices": [{"index": index, "delta": delta, "finish_reason": finish}],
+        }
+
+    if scenario in ("after_role", "empty", "new_choice"):
+        chunks.append(chunk({"role": "assistant", **({"tool_calls": []} if scenario == "empty" else {})}, index=0))
+
+    other_tool = {"index": 0, "id": "call_other", "type": "function", "function": {"name": "other", "arguments": "{}"}}
+    if scenario == "existing_tool":
+        chunks.append(chunk({"role": "assistant", "tool_calls": [other_tool]}))
+
+    # Match the report: the name and first argument fragment share an index
+    # within one chunk. Also split a later argument delta into two entries.
+    first_tools = [
+        {"index": tool_index, "id": "call_files", "type": "function", "function": {"name": "list_files"}},
+        {"index": tool_index, "function": {"arguments": '{"path"'}},
+    ]
+    if scenario == "parallel":
+        first_tools.insert(0, other_tool)
+    chunks.extend(
+        [
+            chunk({"role": "assistant", "tool_calls": first_tools}),
+            chunk(
+                {
+                    "tool_calls": [
+                        {"index": tool_index, "function": {"arguments": ": "}},
+                        {"index": tool_index, "function": {"arguments": '"."}'}},
+                    ]
+                }
+            ),
+            chunk({}, finish="tool_calls"),
+        ]
+    )
+    return chunks
+
+
+def mock_tool_stream(router: MockRouter, chunks: list[dict[str, Any]]) -> None:
+    body = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + "data: [DONE]\n\n"
+    router.post("/chat/completions").mock(
+        return_value=httpx2.Response(200, text=body, headers={"content-type": "text/event-stream"})
+    )
+
+
+def assert_tool_stream(
+    completion: ParsedChatCompletion[None],
+    events: list[ChatCompletionStreamEvent[None]],
+    chunks: list[dict[str, Any]],
+) -> None:
+    calls = completion.choices[-1].message.tool_calls
+    assert calls is not None
+    expected_count = (
+        2
+        if any(
+            tool.get("id") == "call_other"
+            for chunk in chunks
+            for choice in chunk["choices"]
+            for tool in choice["delta"].get("tool_calls", [])
+        )
+        else 1
+    )
+    assert len(calls) == expected_count
+    assert calls[-1].id == "call_files"
+    assert calls[-1].function.name == "list_files"
+    assert calls[-1].function.arguments == '{"path": "."}'
+    assert calls[-1].function.parsed_arguments == ListFilesArgs(path=".")
+    if expected_count == 2:
+        assert calls[0].id == "call_other"
+        assert calls[0].function.arguments == "{}"
+
+    deltas = [
+        event for event in events if event.type == "tool_calls.function.arguments.delta" and event.name == "list_files"
+    ]
+    assert [event.arguments_delta for event in deltas] == ["", '{"path"', ": ", '"."}']
+    done = [
+        event for event in events if event.type == "tool_calls.function.arguments.done" and event.name == "list_files"
+    ]
+    assert len(done) == 1
+    assert done[0].arguments == '{"path": "."}'
+    assert done[0].parsed_arguments == ListFilesArgs(path=".")
+    # Snapshot normalization must not change the raw chunks exposed to callers.
+    assert [event.chunk.to_dict() for event in events if event.type == "chunk"] == chunks
+
+
+@pytest.mark.respx2(base_url=base_url)
+def test_duplicate_tool_call_deltas(
+    client: OpenAI, respx2_mock: MockRouter, tool_call_chunks: list[dict[str, Any]]
+) -> None:
+    mock_tool_stream(respx2_mock, tool_call_chunks)
+    with client.chat.completions.stream(
+        model="gpt-4o-mini", messages=[], tools=[openai.pydantic_function_tool(ListFilesArgs, name="list_files")]
+    ) as stream:
+        events = list(stream)
+        assert_tool_stream(stream.get_final_completion(), events, tool_call_chunks)
+
+
+@pytest.mark.respx2(base_url=base_url)
+async def test_async_duplicate_tool_call_deltas(
+    async_client: AsyncOpenAI, respx2_mock: MockRouter, tool_call_chunks: list[dict[str, Any]]
+) -> None:
+    mock_tool_stream(respx2_mock, tool_call_chunks)
+    async with async_client.chat.completions.stream(
+        model="gpt-4o-mini", messages=[], tools=[openai.pydantic_function_tool(ListFilesArgs, name="list_files")]
+    ) as stream:
+        events = [event async for event in stream]
+        assert_tool_stream(await stream.get_final_completion(), events, tool_call_chunks)

tests/lib/test_assistant_tool_call_deltas.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import json
+
+import httpx2
+import pytest
+
+from openai import OpenAI, AsyncOpenAI
+from tests.respx2 import MockRouter
+from openai.types.beta.threads.runs import RunStep
+
+from ..conftest import base_url
+
+
+@pytest.fixture(params=[False, True], ids=["empty", "snapshot-without-index"])
+def assistant_tool_stream(request: pytest.FixtureRequest, respx2_mock: MockRouter) -> None:
+    tool = {"id": "call_files", "type": "function", "function": {"name": "list_files", "arguments": ""}}
+    step = {
+        "id": "step_test",
+        "object": "thread.run.step",
+        "assistant_id": "asst_test",
+        "thread_id": "thread_test",
+        "run_id": "run_test",
+        "created_at": 0,
+        "status": "in_progress",
+        "type": "tool_calls",
+        "step_details": {"type": "tool_calls", "tool_calls": [tool] if request.param else []},
+    }
+    fragments = [
+        {"index": 0, "type": "function", "function": {"arguments": '{"'}},
+        {"index": 0, "type": "function", "function": {"arguments": 'path"'}},
+    ]
+    if not request.param:
+        fragments.insert(0, {"index": 0, **tool})
+    events = [
+        ("thread.run.step.created", step),
+        (
+            "thread.run.step.delta",
+            {
+                "id": "step_test",
+                "object": "thread.run.step.delta",
+                "delta": {
+                    "step_details": {"type": "tool_calls", "tool_calls": fragments},
+                },
+            },
+        ),
+        (
+            "thread.run.step.delta",
+            {
+                "id": "step_test",
+                "object": "thread.run.step.delta",
+                "delta": {
+                    "step_details": {
+                        "type": "tool_calls",
+                        "tool_calls": [
+                            {"index": 0, "type": "function", "function": {"arguments": ': "."}'}},
+                        ],
+                    },
+                },
+            },
+        ),
+    ]
+    body = "".join(f"event: {name}\ndata: {json.dumps(data)}\n\n" for name, data in events) + "data: [DONE]\n\n"
+    respx2_mock.post("/threads/thread_test/runs").mock(
+        return_value=httpx2.Response(200, text=body, headers={"content-type": "text/event-stream"})
+    )
+
+
+def assert_tool_arguments(steps: list[RunStep]) -> None:
+    assert len(steps) == 1
+    details = steps[0].step_details
+    assert details.type == "tool_calls"
+    assert len(details.tool_calls) == 1
+    tool = details.tool_calls[0]
+    assert tool.type == "function"
+    assert tool.id == "call_files"
+    assert tool.function.name == "list_files"
+    assert json.loads(tool.function.arguments) == {"path": "."}
+
+
+@pytest.mark.respx2(base_url=base_url)
+@pytest.mark.usefixtures("assistant_tool_stream")
+@pytest.mark.filterwarnings("ignore:The Assistants API is deprecated:DeprecationWarning")
+def test_assistant_tool_call_deltas(client: OpenAI) -> None:
+    with client.beta.threads.runs.stream(thread_id="thread_test", assistant_id="asst_test") as stream:  # pyright: ignore[reportDeprecated]
+        assert_tool_arguments(stream.get_final_run_steps())
+
+
+@pytest.mark.respx2(base_url=base_url)
+@pytest.mark.usefixtures("assistant_tool_stream")
+@pytest.mark.filterwarnings("ignore:The Assistants API is deprecated:DeprecationWarning")
+async def test_async_assistant_tool_call_deltas(async_client: AsyncOpenAI) -> None:
+    async with async_client.beta.threads.runs.stream(thread_id="thread_test", assistant_id="asst_test") as stream:  # pyright: ignore[reportDeprecated]
+        assert_tool_arguments(await stream.get_final_run_steps())