core/agent: parallel tool calls and tool-result char budget

Truncate large tool outputs and run multiple tool calls in one round
concurrently when enabled.

src/plyngent/agent/budget.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+DEFAULT_TOOL_RESULT_MAX_CHARS = 32_000
+
+
+def truncate_tool_result(text: str, max_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS) -> str:
+    """Cap tool output so huge dumps do not flood model context."""
+    if max_chars < 1:
+        return text
+    if len(text) <= max_chars:
+        return text
+    omitted = len(text) - max_chars
+    return f"{text[:max_chars]}\n...[truncated {omitted} characters]"

src/plyngent/agent/loop.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from typing import TYPE_CHECKING, cast
 
 from msgspec import UNSET
@@ -15,6 +16,7 @@ from plyngent.lmproto.openai_compatible.model import (
 )
 from plyngent.typedef import Unset  # noqa: TC001
 
+from .budget import DEFAULT_TOOL_RESULT_MAX_CHARS, truncate_tool_result
 from .events import (
     AgentEvent,
     AssistantMessageEvent,
@@ -39,25 +41,52 @@ if TYPE_CHECKING:
 DEFAULT_MAX_ROUNDS = 32
 
 
-async def _execute_tool_calls(
+async def _run_one_tool(
     tools: ToolRegistry,
-    tool_calls: Sequence[AnyAssistantToolCall],
+    call: AnyAssistantToolCall,
-    messages: list[AnyChatMessage],
+    *,
-) -> AsyncIterator[AgentEvent]:
+    max_result_chars: int,
-    for call in tool_calls:
+) -> tuple[ToolChatMessage, ErrorEvent | None]:
-        yield ToolCallEvent(tool_call=call)
     if isinstance(call, AssistantFunctionToolCall):
         try:
             result_text = await tools.execute(call.function.name, call.function.arguments)
         except Exception as exc:  # noqa: BLE001
             result_text = f"error: tool {call.function.name!r} failed: {exc}"
-                yield ErrorEvent(message=result_text)
+            err = ErrorEvent(message=result_text)
-            tool_msg = ToolChatMessage(content=result_text, tool_call_id=call.id)
+            truncated = truncate_tool_result(result_text, max_result_chars)
-        else:
+            return ToolChatMessage(content=truncated, tool_call_id=call.id), err
-            tool_msg = ToolChatMessage(
+        truncated = truncate_tool_result(result_text, max_result_chars)
+        return ToolChatMessage(content=truncated, tool_call_id=call.id), None
+    msg = ToolChatMessage(
         content="error: custom tool calls are not supported",
         tool_call_id=call.id,
     )
+    return msg, None
+
+
+async def _execute_tool_calls(
+    tools: ToolRegistry,
+    tool_calls: Sequence[AnyAssistantToolCall],
+    messages: list[AnyChatMessage],
+    *,
+    max_result_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS,
+    parallel: bool = True,
+) -> AsyncIterator[AgentEvent]:
+    for call in tool_calls:
+        yield ToolCallEvent(tool_call=call)
+
+    if parallel and len(tool_calls) > 1:
+        results = await asyncio.gather(
+            *[_run_one_tool(tools, call, max_result_chars=max_result_chars) for call in tool_calls]
+        )
+    else:
+        results = [
+            await _run_one_tool(tools, call, max_result_chars=max_result_chars) for call in tool_calls
+        ]
+
+    for tool_msg, err in results:
+        if err is not None:
+            yield err
         messages.append(tool_msg)
         yield ToolResultEvent(message=tool_msg)
 
@@ -127,11 +156,14 @@ async def run_chat_loop(
     temperature: float | None = None,
     on_limit: LimitContinueHook | None = None,
     stream: bool = True,
+    max_tool_result_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS,
+    parallel_tools: bool = True,
 ) -> AsyncIterator[AgentEvent]:
     """Multi-round chat/tool loop; mutates ``messages`` in place and yields events.
 
     When ``stream=True``, uses ``chat_completions(..., stream=True)`` and yields
     text deltas as chunks arrive; tool calls are merged from stream deltas.
+    Multiple tool calls in one round run in parallel when ``parallel_tools``.
     """
     tool_items: Sequence[AnyToolItem] | None = None
     if tools is not None and len(tools) > 0:
@@ -166,7 +198,13 @@ async def run_chat_loop(
                 return
             if tools is None:
                 return
-            async for event in _execute_tool_calls(tools, tool_calls, messages):
+            async for event in _execute_tool_calls(
+                tools,
+                tool_calls,
+                messages,
+                max_result_chars=max_tool_result_chars,
+                parallel=parallel_tools,
+            ):
                 yield event
 
         reason = f"tool loop reached {allowance} rounds (used {rounds_used})"

tests/test_agent/test_budget.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from plyngent.agent.budget import truncate_tool_result
+
+
+def test_truncate_tool_result_short() -> None:
+    assert truncate_tool_result("hello", 100) == "hello"
+
+
+def test_truncate_tool_result_long() -> None:
+    text = "a" * 50
+    out = truncate_tool_result(text, 20)
+    assert out.startswith("a" * 20)
+    assert "truncated" in out
+    assert "30" in out

tests/test_agent/test_parallel_tools.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Literal, overload
+
+from msgspec import UNSET
+
+from plyngent.agent import ChatAgent, ToolRegistry, tool
+from plyngent.lmproto.openai_compatible.model import (
+    AssistantChatMessage,
+    AssistantFunctionTool,
+    AssistantFunctionToolCall,
+    ChatCompletionChoice,
+    ChatCompletionChunk,
+    ChatCompletionResponse,
+    ChatCompletionsParam,
+    ChunkChoice,
+    DeltaMessage,
+    StreamFunctionDelta,
+    StreamToolCallDelta,
+    ToolChatMessage,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+
+def _response(message: AssistantChatMessage) -> ChatCompletionResponse:
+    return ChatCompletionResponse(
+        id="1",
+        object="chat.completion",
+        created=0,
+        model="t",
+        choices=[ChatCompletionChoice(index=0, message=message, logprobs={}, finish_reason="stop")],
+        system_fingerprint="",
+        usage={},
+    )
+
+
+class ParallelScriptedClient:
+    step: int
+
+    def __init__(self) -> None:
+        self.step = 0
+
+    @overload
+    async def chat_completions(
+        self, param: ChatCompletionsParam, *, stream: Literal[False] = False
+    ) -> ChatCompletionResponse: ...
+
+    @overload
+    async def chat_completions(
+        self, param: ChatCompletionsParam, *, stream: Literal[True]
+    ) -> AsyncIterator[ChatCompletionChunk]: ...
+
+    async def chat_completions(
+        self, param: ChatCompletionsParam, *, stream: bool = False
+    ) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
+        del param
+        self.step += 1
+        if self.step == 1:
+            msg = AssistantChatMessage(
+                content="",
+                tool_calls=[
+                    AssistantFunctionToolCall(
+                        id="1",
+                        function=AssistantFunctionTool(name="slow_a", arguments="{}"),
+                    ),
+                    AssistantFunctionToolCall(
+                        id="2",
+                        function=AssistantFunctionTool(name="slow_b", arguments="{}"),
+                    ),
+                ],
+            )
+        else:
+            msg = AssistantChatMessage(content="done")
+        response = _response(msg)
+        if stream:
+            return self._stream(response)
+        return response
+
+    async def _stream(self, response: ChatCompletionResponse) -> AsyncIterator[ChatCompletionChunk]:
+        message = response.choices[0].message
+        if isinstance(message.content, str) and message.content:
+            yield ChatCompletionChunk(
+                id="1",
+                object="chat.completion.chunk",
+                created=0,
+                model="t",
+                choices=[ChunkChoice(index=0, delta=DeltaMessage(content=message.content), finish_reason=None)],
+            )
+        tool_calls = message.tool_calls
+        if tool_calls is not UNSET and tool_calls:
+            deltas: list[StreamToolCallDelta] = []
+            for i, call in enumerate(tool_calls):
+                if isinstance(call, AssistantFunctionToolCall):
+                    deltas.append(
+                        StreamToolCallDelta(
+                            index=i,
+                            id=call.id,
+                            type="function",
+                            function=StreamFunctionDelta(
+                                name=call.function.name,
+                                arguments=call.function.arguments,
+                            ),
+                        )
+                    )
+            yield ChatCompletionChunk(
+                id="1",
+                object="chat.completion.chunk",
+                created=0,
+                model="t",
+                choices=[
+                    ChunkChoice(
+                        index=0,
+                        delta=DeltaMessage(tool_calls=deltas),
+                        finish_reason="tool_calls",
+                    )
+                ],
+            )
+
+
+async def test_parallel_tool_calls_run_concurrently() -> None:
+    started = 0
+    expected_parallel = 2
+    gate = asyncio.Event()
+    wait_timeout = 2.0
+
+    @tool
+    async def slow_a() -> str:
+        nonlocal started
+        started += 1
+        if started >= expected_parallel:
+            _ = gate.set()
+        _ = await asyncio.wait_for(gate.wait(), timeout=wait_timeout)
+        return "a"
+
+    @tool
+    async def slow_b() -> str:
+        nonlocal started
+        started += 1
+        if started >= expected_parallel:
+            _ = gate.set()
+        _ = await asyncio.wait_for(gate.wait(), timeout=wait_timeout)
+        return "b"
+
+    registry = ToolRegistry([slow_a, slow_b])
+    agent = ChatAgent(
+        ParallelScriptedClient(),
+        model="m",
+        tools=registry,
+        stream=True,
+        parallel_tools=True,
+    )
+    events = [e async for e in agent.run("go")]
+    tool_msgs = [e for e in agent.messages if isinstance(e, ToolChatMessage)]
+    assert {m.content for m in tool_msgs} == {"a", "b"}
+    assert any(getattr(e, "content", None) == "done" for e in events)