diff --git a/src/plyngent/agent/budget.py b/src/plyngent/agent/budget.py index 035a7c6..0397200 100644 --- a/src/plyngent/agent/budget.py +++ b/src/plyngent/agent/budget.py @@ -1,6 +1,24 @@ from __future__ import annotations +from typing import TYPE_CHECKING + +import msgspec +from msgspec import UNSET + +from plyngent.lmproto.openai_compatible.model import ( + AssistantChatMessage, + ToolChatMessage, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + + from plyngent.lmproto.openai_compatible.model import AnyChatMessage + DEFAULT_TOOL_RESULT_MAX_CHARS = 32_000 +DEFAULT_CONTEXT_MAX_CHARS = 200_000 +DEFAULT_OLD_TOOL_RESULT_CHARS = 800 +DEFAULT_RECENT_TOOL_RESULTS = 4 def truncate_tool_result(text: str, max_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS) -> str: @@ -11,3 +29,117 @@ def truncate_tool_result(text: str, max_chars: int = DEFAULT_TOOL_RESULT_MAX_CHA return text omitted = len(text) - max_chars return f"{text[:max_chars]}\n...[truncated {omitted} characters]" + + +def estimate_message_chars(message: AnyChatMessage) -> int: + """Rough character estimate for soft context budgeting (not tokenizer-accurate).""" + total = 0 + content = getattr(message, "content", None) + if isinstance(content, str): + total += len(content) + if isinstance(message, AssistantChatMessage): + tool_calls = message.tool_calls + if tool_calls is not UNSET and tool_calls: + for call in tool_calls: + total += len(getattr(call, "id", "") or "") + function = getattr(call, "function", None) + if function is not None: + total += len(getattr(function, "name", "") or "") + total += len(getattr(function, "arguments", "") or "") + custom = getattr(call, "custom", None) + if custom is not None: + total += len(getattr(custom, "name", "") or "") + total += len(getattr(custom, "input", "") or "") + if isinstance(message, ToolChatMessage): + total += len(message.tool_call_id) + return total + + +def estimate_messages_chars(messages: Sequence[AnyChatMessage]) -> int: + return sum(estimate_message_chars(m) for m in messages) + + +def _shrink_tool(message: ToolChatMessage, max_chars: int) -> ToolChatMessage: + if len(message.content) <= max_chars: + return message + return msgspec.structs.replace( + message, + content=truncate_tool_result(message.content, max_chars), + ) + + +def _tool_indices(messages: Sequence[AnyChatMessage]) -> list[int]: + return [i for i, m in enumerate(messages) if isinstance(m, ToolChatMessage)] + + +def _protect_indices(tool_indices: Sequence[int], keep_recent: int) -> set[int]: + if keep_recent <= 0: + return set() + return set(tool_indices[-keep_recent:]) + + +def _shrink_except( + messages: list[AnyChatMessage], + tool_indices: Sequence[int], + protect: set[int], + max_chars: int, +) -> None: + for idx in tool_indices: + if idx in protect: + continue + tool_msg = messages[idx] + if isinstance(tool_msg, ToolChatMessage): + messages[idx] = _shrink_tool(tool_msg, max_chars) + + +def _shrink_largest( + messages: list[AnyChatMessage], + tool_indices: Sequence[int], + *, + max_chars: int, + shrink_cap: int, +) -> None: + def tool_len(i: int) -> int: + msg = messages[i] + return len(msg.content) if isinstance(msg, ToolChatMessage) else 0 + + for idx in sorted(tool_indices, key=tool_len, reverse=True): + if estimate_messages_chars(messages) <= max_chars: + return + tool_msg = messages[idx] + if isinstance(tool_msg, ToolChatMessage): + messages[idx] = _shrink_tool(tool_msg, shrink_cap) + + +def compact_messages_for_request( + messages: Sequence[AnyChatMessage], + *, + max_chars: int = DEFAULT_CONTEXT_MAX_CHARS, + old_tool_result_chars: int = DEFAULT_OLD_TOOL_RESULT_CHARS, + keep_recent_tool_results: int = DEFAULT_RECENT_TOOL_RESULTS, +) -> list[AnyChatMessage]: + """Return a request-time copy with older tool dumps shrunk if over budget. + + Does not mutate the original history (full results stay for persistence/UI). + ``max_chars < 1`` disables compacting. + """ + out: list[AnyChatMessage] = list(messages) + if max_chars < 1 or estimate_messages_chars(out) <= max_chars: + return out + + indices = _tool_indices(out) + if not indices: + return out + + protect = _protect_indices(indices, keep_recent_tool_results) + _shrink_except(out, indices, protect, old_tool_result_chars) + if estimate_messages_chars(out) <= max_chars: + return out + + _shrink_largest( + out, + indices, + max_chars=max_chars, + shrink_cap=max(64, old_tool_result_chars // 2), + ) + return out diff --git a/src/plyngent/agent/loop.py b/src/plyngent/agent/loop.py index d1055fc..5a913d6 100644 --- a/src/plyngent/agent/loop.py +++ b/src/plyngent/agent/loop.py @@ -16,7 +16,12 @@ 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 .budget import ( + DEFAULT_CONTEXT_MAX_CHARS, + DEFAULT_TOOL_RESULT_MAX_CHARS, + compact_messages_for_request, + truncate_tool_result, +) from .events import ( AgentEvent, AssistantMessageEvent, @@ -41,18 +46,28 @@ if TYPE_CHECKING: DEFAULT_MAX_ROUNDS = 32 +def _raise_if_cancelled() -> None: + """Cooperative cancel points between model/tool steps.""" + task = asyncio.current_task() + if task is not None and task.cancelling(): + raise asyncio.CancelledError + + async def _run_one_tool( tools: ToolRegistry, call: AnyAssistantToolCall, *, max_result_chars: int, ) -> tuple[ToolChatMessage, ErrorEvent | None]: + _raise_if_cancelled() if isinstance(call, AssistantFunctionToolCall): try: result_text = await tools.execute(call.function.name, call.function.arguments) + except asyncio.CancelledError: + raise except Exception as exc: # noqa: BLE001 result_text = f"error: tool {call.function.name!r} failed: {exc}" - err = ErrorEvent(message=result_text) + err = ErrorEvent(message=result_text, retryable=True, source="tool") 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) @@ -61,7 +76,11 @@ async def _run_one_tool( content="error: custom tool calls are not supported", tool_call_id=call.id, ) - return msg, None + return msg, ErrorEvent( + message=msg.content, + retryable=False, + source="tool", + ) async def _execute_tool_calls( @@ -75,6 +94,7 @@ async def _execute_tool_calls( for call in tool_calls: yield ToolCallEvent(tool_call=call) + _raise_if_cancelled() 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] @@ -158,12 +178,14 @@ async def run_chat_loop( stream: bool = True, max_tool_result_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS, parallel_tools: bool = True, + max_context_chars: int = DEFAULT_CONTEXT_MAX_CHARS, ) -> 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``. + Request payloads may shrink older tool results when over ``max_context_chars``. """ tool_items: Sequence[AnyToolItem] | None = None if tools is not None and len(tools) > 0: @@ -175,8 +197,13 @@ async def run_chat_loop( while True: while rounds_used < allowance: rounds_used += 1 + _raise_if_cancelled() + request_messages = compact_messages_for_request( + messages, + max_chars=max_context_chars, + ) param = ChatCompletionsParam( - messages=list(messages), + messages=request_messages, model=model, temperature=temperature if temperature is not None else UNSET, tools=list(tool_items) if tool_items is not None else UNSET, diff --git a/tests/test_agent/test_budget.py b/tests/test_agent/test_budget.py index 2640283..0f2d758 100644 --- a/tests/test_agent/test_budget.py +++ b/tests/test_agent/test_budget.py @@ -1,6 +1,7 @@ from __future__ import annotations -from plyngent.agent.budget import truncate_tool_result +from plyngent.agent.budget import estimate_message_chars, truncate_tool_result +from plyngent.lmproto.openai_compatible.model import ToolChatMessage, UserChatMessage def test_truncate_tool_result_short() -> None: @@ -13,3 +14,9 @@ def test_truncate_tool_result_long() -> None: assert out.startswith("a" * 20) assert "truncated" in out assert "30" in out + + +def test_estimate_message_chars() -> None: + assert estimate_message_chars(UserChatMessage(content="hello")) == 5 # noqa: PLR2004 + tool = ToolChatMessage(content="abc", tool_call_id="id1") + assert estimate_message_chars(tool) == 6 # noqa: PLR2004 diff --git a/tests/test_agent/test_compact.py b/tests/test_agent/test_compact.py new file mode 100644 index 0000000..d152779 --- /dev/null +++ b/tests/test_agent/test_compact.py @@ -0,0 +1,224 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal, overload + +from msgspec import UNSET + +from plyngent.agent.budget import ( + compact_messages_for_request, + estimate_messages_chars, + truncate_tool_result, +) +from plyngent.agent.loop import run_chat_loop +from plyngent.lmproto.openai_compatible.model import ( + AssistantChatMessage, + AssistantFunctionTool, + AssistantFunctionToolCall, + ChatCompletionChoice, + ChatCompletionChunk, + ChatCompletionResponse, + ChatCompletionsParam, + ChunkChoice, + DeltaMessage, + StreamFunctionDelta, + StreamToolCallDelta, + ToolChatMessage, + UserChatMessage, +) + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from plyngent.lmproto.openai_compatible.model import AnyChatMessage + + +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 + + +def test_compact_shrinks_old_tool_results() -> None: + messages: list[AnyChatMessage] = [ + UserChatMessage(content="start"), + AssistantChatMessage( + content="", + tool_calls=[ + AssistantFunctionToolCall( + id="1", + function=AssistantFunctionTool(name="t", arguments="{}"), + ) + ], + ), + ToolChatMessage(content="OLD" * 200, tool_call_id="1"), + UserChatMessage(content="again"), + AssistantChatMessage( + content="", + tool_calls=[ + AssistantFunctionToolCall( + id="2", + function=AssistantFunctionTool(name="t", arguments="{}"), + ) + ], + ), + ToolChatMessage(content="NEW" * 50, tool_call_id="2"), + ] + original_old = messages[2] + assert isinstance(original_old, ToolChatMessage) + original_len = len(original_old.content) + + compacted = compact_messages_for_request( + messages, + max_chars=estimate_messages_chars(messages) - 1, + old_tool_result_chars=40, + keep_recent_tool_results=1, + ) + assert isinstance(compacted[2], ToolChatMessage) + assert len(compacted[2].content) < original_len + assert "truncated" in compacted[2].content + # Full history unchanged + assert isinstance(messages[2], ToolChatMessage) + assert len(messages[2].content) == original_len + # Recent tool kept + assert isinstance(compacted[5], ToolChatMessage) + assert compacted[5].content == "NEW" * 50 + + +def test_compact_disabled_when_max_chars_zero() -> None: + messages: list[AnyChatMessage] = [ + ToolChatMessage(content="x" * 500, tool_call_id="1"), + ] + out = compact_messages_for_request(messages, max_chars=0) + assert out[0] is messages[0] or ( + isinstance(out[0], ToolChatMessage) and out[0].content == "x" * 500 + ) + + +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 CaptureClient: + _responses: list[ChatCompletionResponse] + calls: list[ChatCompletionsParam] + + def __init__(self, responses: list[ChatCompletionResponse]) -> None: + self._responses = list(responses) + self.calls = [] + + @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]: + self.calls.append(param) + response = self._responses.pop(0) + if stream: + + async def as_stream() -> 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", + ) + ], + ) + + return as_stream() + return response + + +async def test_loop_sends_compacted_request_not_history() -> None: + big = "Z" * 500 + history: list[AnyChatMessage] = [ + UserChatMessage(content="u"), + AssistantChatMessage( + content="", + tool_calls=[ + AssistantFunctionToolCall( + id="1", + function=AssistantFunctionTool(name="t", arguments="{}"), + ) + ], + ), + ToolChatMessage(content=big, tool_call_id="1"), + UserChatMessage(content="next"), + ] + client = CaptureClient([_response(AssistantChatMessage(content="ok"))]) + _ = [ + e + async for e in run_chat_loop( + client, + history, + model="m", + stream=False, + max_context_chars=200, + max_tool_result_chars=50, + ) + ] + assert client.calls + sent = client.calls[0].messages + tool_sent = next(m for m in sent if isinstance(m, ToolChatMessage)) + assert len(tool_sent.content) < len(big) + # In-memory history still full for the old tool result + assert isinstance(history[2], ToolChatMessage) + assert history[2].content == big