mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
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.
This commit is contained in:
@@ -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]"
|
||||
+51
-13
@@ -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 _run_one_tool(
|
||||
tools: ToolRegistry,
|
||||
call: AnyAssistantToolCall,
|
||||
*,
|
||||
max_result_chars: int,
|
||||
) -> tuple[ToolChatMessage, ErrorEvent | None]:
|
||||
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}"
|
||||
err = ErrorEvent(message=result_text)
|
||||
truncated = truncate_tool_result(result_text, max_result_chars)
|
||||
return ToolChatMessage(content=truncated, tool_call_id=call.id), err
|
||||
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 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)
|
||||
tool_msg = ToolChatMessage(content=result_text, tool_call_id=call.id)
|
||||
else:
|
||||
tool_msg = ToolChatMessage(
|
||||
content="error: custom tool calls are not supported",
|
||||
tool_call_id=call.id,
|
||||
)
|
||||
|
||||
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})"
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user