diff --git a/src/plyngent/agent/budget.py b/src/plyngent/agent/budget.py new file mode 100644 index 0000000..035a7c6 --- /dev/null +++ b/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]" diff --git a/src/plyngent/agent/loop.py b/src/plyngent/agent/loop.py index 617f439..d1055fc 100644 --- a/src/plyngent/agent/loop.py +++ b/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 _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})" diff --git a/tests/test_agent/test_budget.py b/tests/test_agent/test_budget.py new file mode 100644 index 0000000..2640283 --- /dev/null +++ b/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 diff --git a/tests/test_agent/test_parallel_tools.py b/tests/test_agent/test_parallel_tools.py new file mode 100644 index 0000000..f7e0dfc --- /dev/null +++ b/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)