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