fix: fix streaming of function/tool call arguments (#239)

Author: adubovik

aidial_sdk/chat_completion/chunks.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, TypedDict
+from typing import Any, Dict, List, Literal, Optional, TypedDict
 
 from aidial_sdk.chat_completion.enums import FinishReason, Status
 from aidial_sdk.exceptions import HTTPException as DIALException
@@ -100,6 +100,7 @@ class FunctionToolCallChunk(BaseChunk):
     choice_index: int
     call_index: int
     id: Optional[str]
+    type: Optional[Literal["function"]]
     name: Optional[str]
     arguments: Optional[str]
 
@@ -108,12 +109,14 @@ def __init__(
         choice_index: int,
         call_index: int,
         id: Optional[str],
+        type: Optional[Literal["function"]],
         name: Optional[str],
         arguments: Optional[str],
     ):
         self.choice_index = choice_index
         self.call_index = call_index
         self.id = id
+        self.type = type
         self.name = name
         self.arguments = arguments
 
@@ -122,14 +125,15 @@ def to_dict(self):
             "choices": [
                 {
                     "index": self.choice_index,
+                    "finish_reason": None,
                     "delta": {
                         "content": None,
                         "tool_calls": [
                             remove_nones(
                                 {
                                     "index": self.call_index,
                                     "id": self.id,
-                                    "type": "function",
+                                    "type": self.type,
                                     "function": remove_nones(
                                         {
                                             "name": self.name,
@@ -166,6 +170,7 @@ def to_dict(self):
             "choices": [
                 {
                     "index": self.choice_index,
+                    "finish_reason": None,
                     "delta": {
                         "content": None,
                         "function_call": remove_nones(

aidial_sdk/chat_completion/function_call.py

@@ -15,13 +15,17 @@ def __init__(self, choice: ChoiceBase):
     def create_and_send(
         cls, choice: ChoiceBase, name: str, arguments: Optional[str]
     ) -> "FunctionCall":
-        return cls(choice)._send_function_call(name=name, arguments=arguments)
+        return cls(choice)._send_function_call(
+            create=True, name=name, arguments=arguments
+        )
 
     def append_arguments(self, arguments: str) -> "FunctionCall":
-        return self._send_function_call(name=None, arguments=arguments)
+        return self._send_function_call(
+            create=False, name=None, arguments=arguments
+        )
 
     def _send_function_call(
-        self, name: Optional[str], arguments: Optional[str]
+        self, *, create: bool, name: Optional[str], arguments: Optional[str]
     ) -> "FunctionCall":
         if not self._choice.opened:
             raise runtime_error(
@@ -31,7 +35,7 @@ def _send_function_call(
             raise runtime_error(
                 "Trying to add function call to a closed choice"
             )
-        if self._choice.has_function_call:
+        if create and self._choice.has_function_call:
             raise runtime_error(
                 "Trying to add function call to a choice which already has a function call"
             )

aidial_sdk/chat_completion/function_tool_call.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Literal, Optional
 
 from aidial_sdk.chat_completion.choice_base import ChoiceBase
 from aidial_sdk.chat_completion.chunks import FunctionToolCallChunk
@@ -23,14 +23,21 @@ def create_and_send(
         arguments: Optional[str],
     ) -> "FunctionToolCall":
         return cls(choice, index)._send_tool_call(
-            id=id, name=name, arguments=arguments
+            id=id, type="function", name=name, arguments=arguments
         )
 
     def append_arguments(self, arguments: str) -> "FunctionToolCall":
-        return self._send_tool_call(id=None, name=None, arguments=arguments)
+        return self._send_tool_call(
+            id=None, type=None, name=None, arguments=arguments
+        )
 
     def _send_tool_call(
-        self, id: Optional[str], name: Optional[str], arguments: Optional[str]
+        self,
+        *,
+        id: Optional[str],
+        type: Optional[Literal["function"]],
+        name: Optional[str],
+        arguments: Optional[str]
     ) -> "FunctionToolCall":
         if not self._choice.opened:
             raise runtime_error("Trying to add tool call to an unopened choice")
@@ -41,6 +48,7 @@ def _send_tool_call(
             FunctionToolCallChunk(
                 self._choice.index,
                 self._index,
+                type=type,
                 id=id,
                 name=name,
                 arguments=arguments,

tests/__init__.py

@@ -0,0 +1,3 @@
+import pytest
+
+pytest.register_assert_rewrite("tests.utils.chunks")

tests/test_function_calling.py

@@ -0,0 +1,54 @@
+from aidial_sdk.chat_completion import ChatCompletion, Request, Response
+from tests.utils.chunks import (
+    check_sse_stream,
+    create_function_call_chunk,
+    create_single_choice_chunk,
+)
+from tests.utils.client import create_app_client
+
+
+class FunctionCaller(ChatCompletion):
+    async def chat_completion(
+        self, request: Request, response: Response
+    ) -> None:
+        response.set_response_id("test_id")
+        response.set_created(0)
+
+        with response.create_single_choice() as choice:
+            choice.append_content("Test content")
+
+            function_call = choice.create_function_call("function_name")
+            function_call.append_arguments('{"key')
+            function_call.append_arguments('":"')
+            function_call.append_arguments('val"}')
+
+
+def test_function_call_non_streaming():
+    response = create_app_client(FunctionCaller()).post(
+        "chat/completions", json={"messages": [], "stream": False}
+    )
+
+    body = response.json()
+    assert body["choices"][0]["message"]["function_call"] == {
+        "name": "function_name",
+        "arguments": '{"key":"val"}',
+    }
+
+
+def test_function_call_streaming():
+    response = create_app_client(FunctionCaller()).post(
+        "chat/completions", json={"messages": [], "stream": True}
+    )
+
+    check_sse_stream(
+        response.iter_lines(),
+        [
+            create_single_choice_chunk({"role": "assistant"}),
+            create_single_choice_chunk({"content": "Test content"}),
+            create_function_call_chunk(name="function_name"),
+            create_function_call_chunk(arguments='{"key'),
+            create_function_call_chunk(arguments='":"'),
+            create_function_call_chunk(arguments='val"}'),
+            create_single_choice_chunk({}, "function_call"),
+        ],
+    )

tests/test_tool_calling.py

@@ -0,0 +1,83 @@
+from aidial_sdk.chat_completion import ChatCompletion, Request, Response
+from tests.utils.chunks import (
+    check_sse_stream,
+    create_single_choice_chunk,
+    create_tool_call_chunk,
+)
+from tests.utils.client import create_app_client
+
+
+class ToolCaller(ChatCompletion):
+    async def chat_completion(
+        self, request: Request, response: Response
+    ) -> None:
+        response.set_response_id("test_id")
+        response.set_created(0)
+
+        with response.create_single_choice() as choice:
+            choice.append_content("Test content")
+
+            tool_call1 = choice.create_function_tool_call(
+                "tool_call_id1", "tool_name"
+            )
+            tool_call1.append_arguments('{"key')
+            tool_call1.append_arguments('":"')
+            tool_call1.append_arguments('val"}')
+
+            choice.create_function_tool_call(
+                "tool_call_id2", "tool_name", '{"foo":"bar"}'
+            )
+
+
+def test_tool_call_non_streaming():
+    response = create_app_client(ToolCaller()).post(
+        "chat/completions", json={"messages": [], "stream": False}
+    )
+
+    body = response.json()
+    assert body["choices"][0]["message"]["tool_calls"] == [
+        {
+            "id": "tool_call_id1",
+            "type": "function",
+            "function": {
+                "name": "tool_name",
+                "arguments": '{"key":"val"}',
+            },
+        },
+        {
+            "id": "tool_call_id2",
+            "type": "function",
+            "function": {
+                "name": "tool_name",
+                "arguments": '{"foo":"bar"}',
+            },
+        },
+    ]
+
+
+def test_tool_call_streaming():
+    response = create_app_client(ToolCaller()).post(
+        "chat/completions", json={"messages": [], "stream": True}
+    )
+
+    check_sse_stream(
+        response.iter_lines(),
+        [
+            create_single_choice_chunk({"role": "assistant"}),
+            create_single_choice_chunk({"content": "Test content"}),
+            create_tool_call_chunk(
+                0, type="function", id="tool_call_id1", name="tool_name"
+            ),
+            create_tool_call_chunk(0, arguments='{"key'),
+            create_tool_call_chunk(0, arguments='":"'),
+            create_tool_call_chunk(0, arguments='val"}'),
+            create_tool_call_chunk(
+                1,
+                type="function",
+                id="tool_call_id2",
+                name="tool_name",
+                arguments='{"foo":"bar"}',
+            ),
+            create_single_choice_chunk({}, "tool_calls"),
+        ],
+    )

tests/utils/chunks.py

@@ -2,26 +2,29 @@
 import json
 from typing import Iterable, Literal, Optional, Union
 
+from aidial_sdk.utils.json import remove_nones
+
 
 def create_chunk(
     *,
     choice_idx: int = 0,
     delta: dict = {},
     finish_reason: Optional[str] = None,
+    **kwargs,
 ):
     return {
-        "id": "chatcmpl-AQws8iVykPBIQJfnmCQnMEkTLLUUA",
+        "id": "test_id",
         "object": "chat.completion.chunk",
-        "created": 1730986196,
-        "model": "gpt-4o-2024-05-13",
-        "system_fingerprint": "fp_67802d9a6d",
+        "created": 0,
         "choices": [
             {
                 "index": choice_idx,
                 "delta": delta,
                 "finish_reason": finish_reason,
             }
         ],
+        "usage": None,
+        **kwargs,
     }
 
 
@@ -54,23 +57,41 @@ def create_tool_call_chunk(
 ):
     return create_chunk(
         delta={
+            "content": None,
             "tool_calls": [
-                {
-                    "index": idx,
-                    "id": id,
-                    "type": type,
-                    "function": {"name": name, "arguments": arguments},
-                }
-            ]
+                remove_nones(
+                    {
+                        "index": idx,
+                        "id": id,
+                        "type": type,
+                        "function": remove_nones(
+                            {"name": name, "arguments": arguments}
+                        ),
+                    }
+                )
+            ],
+        }
+    )
+
+
+def create_function_call_chunk(
+    *,
+    name: Optional[str] = None,
+    arguments: Optional[str] = None,
+):
+    return create_chunk(
+        delta={
+            "content": None,
+            "function_call": remove_nones(
+                {"name": name, "arguments": arguments}
+            ),
         }
     )
 
 
 def _check_sse_line(actual: str, expected: Union[str, dict]):
     if isinstance(expected, str):
-        assert (
-            actual == expected
-        ), f"actual line != expected line: {actual!r} != {expected!r}"
+        assert actual == expected
         return
 
     assert actual.startswith("data: "), f"Invalid data SSE entry: {actual!r}"
@@ -80,9 +101,8 @@ def _check_sse_line(actual: str, expected: Union[str, dict]):
         actual_dict = json.loads(actual)
     except json.JSONDecodeError:
         raise AssertionError(f"Invalid JSON in data SSE entry: {actual!r}")
-    assert (
-        actual_dict == expected
-    ), f"actual json != expected json: {actual_dict!r} != {expected!r}"
+
+    assert actual_dict == expected
 
 
 ExpectedSSEStream = Iterable[Union[str, dict]]