ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter2/local_llm_serving/test_vllm_structured_streaming.py

@@ -0,0 +1,102 @@
+"""Focused tests for fragmented structured tool calls in VLLMToolAgent.chat_stream."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+from agent import VLLMToolAgent
+
+
+def _chunk(content=None, tool_calls=None):
+    delta = SimpleNamespace(content=content, tool_calls=tool_calls or [])
+    return SimpleNamespace(choices=[SimpleNamespace(delta=delta)])
+
+
+def _fragment(index, *, call_id=None, name=None, arguments=None):
+    return SimpleNamespace(
+        index=index,
+        id=call_id,
+        type="function" if call_id else None,
+        function=SimpleNamespace(name=name, arguments=arguments),
+    )
+
+
+def _agent_with_streams(*streams):
+    agent = VLLMToolAgent.__new__(VLLMToolAgent)
+    agent.conversation_history = []
+    agent.tool_registry = MagicMock()
+    agent.tool_registry.get_tool_schemas.return_value = []
+    agent._format_system_prompt_with_tools = MagicMock(return_value="system")
+    agent.client = MagicMock()
+    agent.client.chat.completions.create.side_effect = [iter(stream) for stream in streams]
+    return agent
+
+
+def test_stream_assembles_fragmented_parallel_tool_calls():
+    agent = _agent_with_streams(
+        [
+            _chunk(tool_calls=[
+                _fragment(0, call_id="call_weather", name="get_", arguments='{"city":'),
+                _fragment(1, call_id="call_time", name="get_time", arguments="{"),
+            ]),
+            _chunk(tool_calls=[
+                _fragment(0, name="weather", arguments='"Paris"}'),
+                _fragment(1, arguments="}"),
+            ]),
+        ],
+        [_chunk(content="Done")],
+    )
+    agent._execute_single_tool = MagicMock(
+        side_effect=lambda call: (f'{call["name"]} result', False)
+    )
+
+    events = list(agent.chat_stream("Use both tools"))
+
+    assert events == [
+        {"type": "tool_call", "content": {"name": "get_weather", "arguments": {"city": "Paris"}}},
+        {"type": "tool_call", "content": {"name": "get_time", "arguments": {}}},
+        {"type": "tool_result", "content": "get_weather result"},
+        {"type": "tool_result", "content": "get_time result"},
+        {"type": "content", "content": "Done"},
+    ]
+    assert agent._execute_single_tool.call_count == 2
+
+    create_calls = agent.client.chat.completions.create.call_args_list
+    assert [call.kwargs["model"] for call in create_calls] == [
+        "Qwen/Qwen3-0.6B",
+        "Qwen/Qwen3-0.6B",
+    ]
+    second_turn_messages = create_calls[1].kwargs["messages"]
+    assistant_message = next(message for message in second_turn_messages if message.get("tool_calls"))
+    assert assistant_message["tool_calls"] == [
+        {
+            "id": "call_weather",
+            "type": "function",
+            "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'},
+        },
+        {
+            "id": "call_time",
+            "type": "function",
+            "function": {"name": "get_time", "arguments": "{}"},
+        },
+    ]
+
+
+def test_stream_reports_malformed_arguments_and_continues():
+    agent = _agent_with_streams(
+        [_chunk(tool_calls=[
+            _fragment(0, call_id="call_bad", name="bad_tool", arguments="{bad")
+        ])],
+        [_chunk(content="Recovered")],
+    )
+    agent._execute_single_tool = MagicMock()
+
+    events = list(agent.chat_stream("Try the tool"))
+
+    assert [event["type"] for event in events] == ["tool_error", "content"]
+    assert events[-1] == {"type": "content", "content": "Recovered"}
+    agent._execute_single_tool.assert_not_called()
+    error_message = next(
+        message for message in agent.conversation_history
+        if message.get("name") == "bad_tool"
+    )
+    assert "Tool call parse exception" in error_message["content"]