fix: stream Work model tool responses

2026-07-26 15:10:56 +08:00
parent 2c17a74203
commit e5207a3dda
2 changed files with 180 additions and 14 deletions
@@ -1,5 +1,6 @@
 from __future__ import annotations

+import json
 from collections.abc import Awaitable, Callable

 import httpx
@@ -78,3 +79,80 @@ async def test_complete_does_not_retry_client_error(settings: Settings) -> None:
        await _complete_with_handler(settings, handler)

    assert calls == 1
+
+
+async def test_complete_aggregates_streamed_tool_call(settings: Settings) -> None:
+    chunks = [
+        {
+            "id": "completion-1",
+            "object": "chat.completion.chunk",
+            "created": 123,
+            "model": "test-model",
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {
+                        "role": "assistant",
+                        "tool_calls": [
+                            {
+                                "index": 0,
+                                "id": "call-1",
+                                "type": "function",
+                                "function": {"name": "ping", "arguments": '{"val'},
+                            }
+                        ],
+                    },
+                    "finish_reason": None,
+                }
+            ],
+        },
+        {
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {
+                        "tool_calls": [
+                            {
+                                "index": 0,
+                                "function": {"arguments": 'ue":"ok"}'},
+                            }
+                        ]
+                    },
+                    "finish_reason": None,
+                }
+            ]
+        },
+        {
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {},
+                    "finish_reason": "tool_calls",
+                }
+            ],
+            "usage": {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7},
+        },
+    ]
+    body = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + "data: [DONE]\n\n"
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        assert json.loads(request.content)["stream"] is True
+        return httpx.Response(
+            200,
+            request=request,
+            headers={"Content-Type": "text/event-stream"},
+            content=body,
+        )
+
+    result = await _complete_with_handler(settings, handler)
+
+    choice = result["choices"][0]
+    assert choice["finish_reason"] == "tool_calls"
+    assert choice["message"]["tool_calls"] == [
+        {
+            "id": "call-1",
+            "type": "function",
+            "function": {"name": "ping", "arguments": '{"value":"ok"}'},
+        }
+    ]
+    assert result["usage"]["total_tokens"] == 7