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

chapter2/local_llm_serving/test_parallel_tools.py

@@ -0,0 +1,249 @@
+"""
+End-to-end test for parallel tool calling with a real LLM (Ollama).
+
+Covers:
+1. Deterministic proof of parallel execution: two 2s-sleep tools must finish
+   in ~2s (not ~4s) through each agent's tool-execution path.
+2. Real-model runs of the book's "Vancouver time + weather" example through:
+   - OllamaNativeAgent.chat / chat_stream (native tool calling)
+   - OllamaOpenAICompatible.chat (OpenAI-compatible endpoint, native tools)
+   - VLLMToolAgent.chat / chat_stream (structured OpenAI-compatible tool calls;
+     Ollama's OpenAI-compatible endpoint stands in for the vLLM server)
+
+Run from this directory:  python3 test_parallel_tools.py
+Requires: ollama serve + ollama pull qwen2.5:7b-instruct-q8_0
+"""
+import json
+import logging
+import time
+from types import SimpleNamespace
+
+from ollama_native import OllamaNativeAgent, OllamaOpenAICompatible
+from agent import VLLMToolAgent
+
+logging.basicConfig(level=logging.WARNING)
+
+SLEEP = 2
+QUERY = "What time is it in Vancouver and what's the current weather in Vancouver?"
+# qwen3:0.6b is the project default but flaky at native tool calling;
+# qwen2.5:7b-instruct-q8_0 (also pulled locally) calls tools reliably.
+REAL_MODEL = "qwen2.5:7b-instruct-q8_0"
+MAX_ATTEMPTS = 4  # small models occasionally skip tool calling; retry
+
+
+# ---------------------------------------------------------------- helpers
+
+def _sleep_tool(name):
+    def fn(tag: str = "") -> dict:
+        time.sleep(SLEEP)
+        return {"tool": name, "slept": SLEEP, "success": True}
+    return fn
+
+
+SLEEP_SCHEMA = {
+    "type": "object",
+    "properties": {"tag": {"type": "string", "description": "optional tag"}},
+    "required": [],
+}
+
+
+def _register_sleep_tools(registry):
+    registry.register_tool("sleep_a", _sleep_tool("sleep_a"), "Sleep tool A", SLEEP_SCHEMA)
+    registry.register_tool("sleep_b", _sleep_tool("sleep_b"), "Sleep tool B", SLEEP_SCHEMA)
+
+
+def _check_parallel(label, elapsed, n=2):
+    status = "PARALLEL OK" if elapsed < SLEEP * 1.8 else "NOT PARALLEL"
+    print(f"  [{label}] {n} x {SLEEP}s tools finished in {elapsed:.1f}s -> {status}")
+    assert elapsed < SLEEP * 1.8, f"{label}: tool calls were not executed in parallel"
+
+
+# ------------------------------------- 1. deterministic parallel checks
+
+def test_native_execute_parallel():
+    print("\n== OllamaNativeAgent._execute_tool_calls (sleep tools) ==")
+    agent = OllamaNativeAgent(model=REAL_MODEL)
+    _register_sleep_tools(agent.tool_registry)
+    calls = [
+        {"function": {"name": "sleep_a", "arguments": {}}},
+        {"function": {"name": "sleep_b", "arguments": {}}},
+    ]
+    start = time.time()
+    results = agent._execute_tool_calls(calls)
+    _check_parallel("native _execute_tool_calls", time.time() - start)
+    assert all(json.loads(r)["success"] for r in results)
+
+
+def _make_chunk(text=None, tool_calls=None):
+    delta = SimpleNamespace(content=text, tool_calls=tool_calls or [])
+    return SimpleNamespace(choices=[SimpleNamespace(delta=delta)])
+
+
+def _tool_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 test_vllm_stream_parallel():
+    print("\n== VLLMToolAgent.chat_stream (fake stream with 2 tool calls) ==")
+    agent = VLLMToolAgent(api_base="http://localhost:11434/v1", api_key="ollama")
+    _register_sleep_tools(agent.tool_registry)
+
+    state = {"n": 0}
+
+    def fake_create(**kwargs):
+        state["n"] += 1
+        if state["n"] == 1:
+            # Split arguments across chunks and interleave two call indexes.
+            return iter([
+                _make_chunk(tool_calls=[
+                    _tool_fragment(
+                        0, call_id="call_a", name="sleep_a", arguments='{"tag":'
+                    ),
+                    _tool_fragment(
+                        1, call_id="call_b", name="sleep_b", arguments='{"tag":'
+                    ),
+                ]),
+                _make_chunk(tool_calls=[
+                    _tool_fragment(0, arguments='"a"}'),
+                    _tool_fragment(1, arguments='"b"}'),
+                ]),
+            ])
+        return iter([_make_chunk("done")])
+
+    agent.client = SimpleNamespace(
+        chat=SimpleNamespace(completions=SimpleNamespace(create=fake_create))
+    )
+
+    start = time.time()
+    events = list(agent.chat_stream("run both sleep tools"))
+    elapsed = time.time() - start
+
+    types = [e["type"] for e in events]
+    print(f"  event types: {types}")
+    assert types.count("tool_call") == 2, f"expected 2 tool_call events: {types}"
+    assert types.count("tool_result") == 2, f"expected 2 tool_result events: {types}"
+    _check_parallel("vllm chat_stream", elapsed)
+
+
+# ------------------------------------- 2. real-model end-to-end runs
+
+def test_native_real_model():
+    print(f"\n== OllamaNativeAgent.chat (real {REAL_MODEL}) ==")
+    agent = OllamaNativeAgent(model=REAL_MODEL)
+    response, tool_msgs = "", []
+    for attempt in range(1, MAX_ATTEMPTS + 1):
+        agent.reset_conversation()
+        response = agent.chat(QUERY)
+        tool_msgs = [m for m in agent.conversation_history if m.get("role") == "tool"]
+        if tool_msgs:
+            break
+        print(f"  attempt {attempt}: model made no tool call, retrying")
+    assistant_tc = [m for m in agent.conversation_history if m.get("tool_calls")]
+    batch = len(assistant_tc[0]["tool_calls"]) if assistant_tc else 0
+    print(f"  tool calls in one turn: {batch} ({'PARALLEL BATCH' if batch > 1 else 'single'})")
+    for m in tool_msgs:
+        print(f"    - {m['content'][:120]}")
+    print(f"  final response: {response[:300]}")
+    assert tool_msgs, "model did not call any tool"
+
+    print(f"\n== OllamaNativeAgent.chat_stream (real {REAL_MODEL}) ==")
+    events = []
+    for attempt in range(1, MAX_ATTEMPTS + 1):
+        agent.reset_conversation()
+        events = list(agent.chat_stream(QUERY))
+        if any(e["type"] == "tool_result" for e in events):
+            break
+        print(f"  attempt {attempt}: no tool call, retrying")
+    tool_calls = [e for e in events if e["type"] == "tool_call"]
+    tool_results = [e for e in events if e["type"] == "tool_result"]
+    print(f"  tool_call events: {len(tool_calls)}, tool_result events: {len(tool_results)}")
+    for e in tool_calls:
+        print(f"    - {e['content']['name']}({e['content']['arguments']})")
+    final = "".join(e["content"] for e in events if e["type"] == "content")
+    print(f"  final content: {final[:300]}")
+    assert tool_results, "streaming path produced no tool results"
+
+
+def test_openai_compat_real_model():
+    print(f"\n== OllamaOpenAICompatible.chat (real {REAL_MODEL}) ==")
+    agent = OllamaOpenAICompatible(model=REAL_MODEL)
+    response, tool_msgs = "", []
+    for attempt in range(1, MAX_ATTEMPTS + 1):
+        agent.reset_conversation()
+        response = agent.chat(QUERY)
+        tool_msgs = [m for m in agent.conversation_history if m.get("role") == "tool"]
+        if tool_msgs:
+            break
+        print(f"  attempt {attempt}: model made no tool call, retrying")
+    print(f"  tool results in history: {len(tool_msgs)}")
+    print(f"  final response: {response[:300]}")
+    assert tool_msgs, "model did not call any tool"
+
+
+def test_vllm_agent_real_model():
+    print(f"\n== VLLMToolAgent.chat (real {REAL_MODEL} via OpenAI endpoint) ==")
+    agent = VLLMToolAgent(api_base="http://localhost:11434/v1", api_key="ollama")
+
+    # Ollama's OpenAI endpoint stands in for the vLLM server. Keep the native
+    # tools payload so both chat paths receive structured tool_calls.
+    real_create = agent.client.chat.completions.create
+
+    def create_structured_mode(**kwargs):
+        kwargs["model"] = REAL_MODEL
+        return real_create(**kwargs)
+
+    agent.client = SimpleNamespace(
+        chat=SimpleNamespace(completions=SimpleNamespace(create=create_structured_mode))
+    )
+
+    response, tool_msgs, batch = "", [], 0
+    for attempt in range(1, MAX_ATTEMPTS + 1):
+        agent.reset_conversation()
+        # temperature 0.7 sampling can degenerate into a tool-call loop with
+        # this handcrafted XML format; 0.3 is stable
+        response = agent.chat(QUERY, temperature=0.3)
+        tool_msgs = [m for m in agent.conversation_history if m.get("name")]
+        assistant_tc = [m for m in agent.conversation_history if m.get("tool_calls")]
+        batch = len(assistant_tc[0]["tool_calls"]) if assistant_tc else 0
+        if tool_msgs and batch <= 10:
+            break
+        print(f"  attempt {attempt}: no tool call or degenerate run ({batch} calls), retrying")
+    print(f"  tool calls in one turn: {batch} ({'PARALLEL BATCH' if batch > 1 else 'single'})")
+    for m in tool_msgs:
+        print(f"    - {m['name']}: {m['content'][:100]}")
+    print(f"  final response: {response[:300]}")
+    assert tool_msgs, "model did not emit a structured tool call"
+
+    print(f"\n== VLLMToolAgent.chat_stream (real {REAL_MODEL}) ==")
+    events = []
+    for attempt in range(1, MAX_ATTEMPTS + 1):
+        agent.reset_conversation()
+        # temperature 0.7 sampling can degenerate into a tool-call loop with
+        # this handcrafted XML format; 0.3 is stable
+        events = list(agent.chat_stream(QUERY, temperature=0.3))
+        n_calls = sum(1 for e in events if e["type"] == "tool_call")
+        if any(e["type"] == "tool_result" for e in events) and n_calls <= 10:
+            break
+        print(f"  attempt {attempt}: no tool call or degenerate run ({n_calls} calls), retrying")
+    tool_calls = [e for e in events if e["type"] == "tool_call"]
+    tool_results = [e for e in events if e["type"] == "tool_result"]
+    print(f"  tool_call events: {len(tool_calls)}, tool_result events: {len(tool_results)}")
+    for e in tool_calls:
+        print(f"    - {e['content']}")
+    final = "".join(e["content"] for e in events if e["type"] == "content")
+    print(f"  final content: {final[:300]}")
+    assert tool_results, "streaming path produced no tool results"
+
+
+if __name__ == "__main__":
+    test_native_execute_parallel()
+    test_vllm_stream_parallel()
+    test_native_real_model()
+    test_openai_compat_real_model()
+    test_vllm_agent_real_model()
+    print("\nALL TESTS PASSED")