mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 07:17:54 +08:00
core/agent: let a tool result carry images beside its text
`ToolOutput` pairs the model-facing text with the images a tool produced. `ToolRegistry.execute` keeps its text-only contract; `execute_output` returns the pair and the loop stores the images on the tool message, so a screenshot tool can answer in pixels instead of a hex placeholder. Only the text is capped by the tool-result limit; the char estimator now counts a present image at its 1024-token ceiling, because the shrinkers never drop one.
This commit is contained in:
@@ -34,6 +34,7 @@ from .tools import ToolBatchConfirmHook as ToolBatchConfirmHook
|
||||
from .tools import ToolConfirmHook as ToolConfirmHook
|
||||
from .tools import ToolConfirmRequest as ToolConfirmRequest
|
||||
from .tools import ToolDefinition as ToolDefinition
|
||||
from .tools import ToolOutput as ToolOutput
|
||||
from .tools import ToolRegistry as ToolRegistry
|
||||
from .tools import ToolTag as ToolTag
|
||||
from .tools import schema_from_callable as schema_from_callable
|
||||
|
||||
@@ -7,7 +7,10 @@ from msgspec import UNSET
|
||||
|
||||
from plyngent.lmproto.openai_compatible.model import (
|
||||
AssistantChatMessage,
|
||||
ChatImagePart,
|
||||
ChatTextPart,
|
||||
ToolChatMessage,
|
||||
UserChatMessage,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -24,6 +27,11 @@ DEFAULT_OLD_TOOL_RESULT_CHARS = 800
|
||||
# Soft-compact: leave this many most-recent tool results at full size.
|
||||
DEFAULT_RECENT_TOOL_RESULTS = 12
|
||||
|
||||
# What one image costs in chars: a provider bills it by its resized size, at most
|
||||
# 1024 tokens, and this estimator works in chars at ~4 chars/token. The shrinkers
|
||||
# below never drop an image, so they must at least count it.
|
||||
IMAGE_CHARS = 4_096
|
||||
|
||||
# Backward-compat alias (older code/docs may still import this name).
|
||||
DEFAULT_CONTEXT_MAX_CHARS = DEFAULT_CONTEXT_MAX_TOKENS * 4
|
||||
|
||||
@@ -50,12 +58,21 @@ def truncate_tool_result(
|
||||
return truncate_generic(text, max_chars)
|
||||
|
||||
|
||||
def _content_chars(content: str | list[ChatTextPart | ChatImagePart]) -> int:
|
||||
if isinstance(content, str):
|
||||
return len(content)
|
||||
return sum(len(part.text) if isinstance(part, ChatTextPart) else IMAGE_CHARS for part in content)
|
||||
|
||||
|
||||
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, UserChatMessage):
|
||||
total += _content_chars(message.content)
|
||||
else:
|
||||
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:
|
||||
@@ -71,6 +88,9 @@ def estimate_message_chars(message: AnyChatMessage) -> int:
|
||||
total += len(getattr(custom, "input", "") or "")
|
||||
if isinstance(message, ToolChatMessage):
|
||||
total += len(message.tool_call_id)
|
||||
images = message.images
|
||||
if images is not UNSET and images:
|
||||
total += IMAGE_CHARS * len(images)
|
||||
return total
|
||||
|
||||
|
||||
|
||||
@@ -93,7 +93,7 @@ async def _run_one_tool(
|
||||
) -> tuple[ToolChatMessage, ErrorEvent | None]:
|
||||
if isinstance(call, AssistantFunctionToolCall):
|
||||
try:
|
||||
result_text = await tools.execute(call.function.name, call.function.arguments, decision=decision)
|
||||
output = await tools.execute_output(call.function.name, call.function.arguments, decision=decision)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc: # noqa: BLE001
|
||||
@@ -101,8 +101,16 @@ async def _run_one_tool(
|
||||
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)
|
||||
return ToolChatMessage(content=truncated, tool_call_id=call.id), None
|
||||
# Only the text is capped; an image the tool returned rides along whole.
|
||||
truncated = truncate_tool_result(output.text, max_result_chars)
|
||||
return (
|
||||
ToolChatMessage(
|
||||
content=truncated,
|
||||
tool_call_id=call.id,
|
||||
images=output.images or UNSET,
|
||||
),
|
||||
None,
|
||||
)
|
||||
msg = ToolChatMessage(
|
||||
content="error: custom tool calls are not supported",
|
||||
tool_call_id=call.id,
|
||||
|
||||
+48
-14
@@ -11,8 +11,9 @@ from pathlib import Path
|
||||
from typing import Any, cast, get_args, get_origin, get_type_hints, overload
|
||||
|
||||
import msgspec
|
||||
from msgspec import Struct
|
||||
|
||||
from plyngent.lmproto.openai_compatible.model import ToolFunction, ToolFunctionItem
|
||||
from plyngent.lmproto.openai_compatible.model import ChatImagePart, ToolFunction, ToolFunctionItem
|
||||
from plyngent.typedef import JSONSchema # noqa: TC001
|
||||
|
||||
type ToolHandler = Callable[..., Any | Awaitable[Any]]
|
||||
@@ -25,6 +26,24 @@ type ToolConfirmHook = Callable[
|
||||
]
|
||||
|
||||
|
||||
class ToolOutput(Struct):
|
||||
"""A tool result: the model-facing text, plus images the model should see.
|
||||
|
||||
Most tools return their text directly; one that produced an image (a
|
||||
screenshot, a chart, a downloaded picture) returns this so the image rides
|
||||
along with the text that explains it. ``execute`` hands the pair to the
|
||||
loop, which keeps the text as the tool message's content and the images
|
||||
beside it.
|
||||
"""
|
||||
|
||||
text: str
|
||||
images: list[ChatImagePart] = msgspec.field(default_factory=list)
|
||||
|
||||
@property
|
||||
def has_images(self) -> bool:
|
||||
return bool(self.images)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ToolConfirmRequest:
|
||||
"""One call of a batch that needs a soft confirm before it may run."""
|
||||
@@ -444,7 +463,7 @@ class ToolRegistry:
|
||||
return f"error: tool {definition.name!r} requires session state (SESSION_STATE) but none is bound"
|
||||
return None
|
||||
|
||||
async def _invoke(self, definition: ToolDefinition, args: dict[str, object]) -> str:
|
||||
async def _invoke(self, definition: ToolDefinition, args: dict[str, object]) -> str | ToolOutput:
|
||||
try:
|
||||
result = definition.handler(**args)
|
||||
if inspect.isawaitable(result):
|
||||
@@ -455,7 +474,7 @@ class ToolRegistry:
|
||||
return f"error: invalid tool arguments: {exc}"
|
||||
except Exception as exc: # noqa: BLE001 — surface tool failures to the model
|
||||
return f"error: tool {definition.name!r} failed: {exc}"
|
||||
if isinstance(result, str):
|
||||
if isinstance(result, str | ToolOutput):
|
||||
return result
|
||||
return msgspec.json.encode(result).decode()
|
||||
|
||||
@@ -645,28 +664,42 @@ class ToolRegistry:
|
||||
*,
|
||||
decision: BatchConfirmDecision | None = None,
|
||||
) -> str:
|
||||
"""Run a tool by name; returns a string result (errors become error text).
|
||||
"""Run a tool by name; returns its text (errors become error text).
|
||||
|
||||
*decision* is the host's answer for this call from
|
||||
:meth:`prepare_batch_confirm`; ``None`` means the call was not part of an
|
||||
aggregated prompt and confirms on its own.
|
||||
An image the tool returned is dropped here — this is the text-only entry
|
||||
point. :meth:`execute_output` returns it; the agent loop uses that one.
|
||||
"""
|
||||
return (await self.execute_output(name, arguments_json, decision=decision)).text
|
||||
|
||||
async def execute_output(
|
||||
self,
|
||||
name: str,
|
||||
arguments_json: str,
|
||||
*,
|
||||
decision: BatchConfirmDecision | None = None,
|
||||
) -> ToolOutput:
|
||||
"""Run a tool by name; returns its text plus any images it produced.
|
||||
|
||||
Errors become error text. *decision* is the host's answer for this call
|
||||
from :meth:`prepare_batch_confirm`; ``None`` means the call was not part
|
||||
of an aggregated prompt and confirms on its own.
|
||||
"""
|
||||
definition = self._tools.get(name)
|
||||
if definition is None:
|
||||
return f"error: unknown tool {name!r}"
|
||||
return ToolOutput(text=f"error: unknown tool {name!r}")
|
||||
args = _decode_tool_args(arguments_json)
|
||||
if isinstance(args, str):
|
||||
return args
|
||||
return ToolOutput(text=args)
|
||||
# Argument shape/coverage problems are answered before any confirm prompt:
|
||||
# the human should never be asked about a call that cannot run.
|
||||
missing_error = _missing_required_error(definition, args)
|
||||
if missing_error is not None:
|
||||
return missing_error
|
||||
return ToolOutput(text=missing_error)
|
||||
|
||||
async def _run() -> str:
|
||||
async def _run() -> ToolOutput:
|
||||
missing = self._check_state_tags(definition)
|
||||
if missing is not None:
|
||||
return missing
|
||||
return ToolOutput(text=missing)
|
||||
# Held before the confirm prompt, so a sibling call writing the same
|
||||
# file waits for this one instead of overtaking its prompt.
|
||||
async with AsyncExitStack() as held:
|
||||
@@ -674,8 +707,9 @@ class ToolRegistry:
|
||||
await held.enter_async_context(lock)
|
||||
denied = await self._maybe_confirm(name, args, tags=definition.tags, decision=decision)
|
||||
if denied is not None:
|
||||
return denied
|
||||
return await self._invoke(definition, args)
|
||||
return ToolOutput(text=denied)
|
||||
result = await self._invoke(definition, args)
|
||||
return result if isinstance(result, ToolOutput) else ToolOutput(text=result)
|
||||
|
||||
# Bind host-provided state for tools that read contextvars (todo, workspace).
|
||||
# Tag enforcement remains gated by ``auto_bind_state``; binding is always useful
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
"""A tool result that carries images beside its text."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from msgspec import UNSET
|
||||
|
||||
from plyngent.agent import ToolOutput, ToolRegistry, tool
|
||||
from plyngent.agent.budget import IMAGE_CHARS, estimate_message_chars
|
||||
from plyngent.agent.loop import _run_one_tool
|
||||
from plyngent.lmproto.openai_compatible.model import (
|
||||
AssistantFunctionTool,
|
||||
AssistantFunctionToolCall,
|
||||
ChatImagePart,
|
||||
ChatImageUrl,
|
||||
ChatTextPart,
|
||||
ToolChatMessage,
|
||||
UserChatMessage,
|
||||
)
|
||||
|
||||
|
||||
def _image() -> ChatImagePart:
|
||||
return ChatImagePart(image_url=ChatImageUrl(url="data:image/png;base64,AAAA", detail="low"))
|
||||
|
||||
|
||||
def _registry() -> ToolRegistry:
|
||||
@tool(register=False)
|
||||
def shot() -> ToolOutput:
|
||||
"""Pretend to take a screenshot."""
|
||||
return ToolOutput(text="image: shot.png (2 KB)", images=[_image()])
|
||||
|
||||
@tool(register=False)
|
||||
def plain() -> str:
|
||||
return "just text"
|
||||
|
||||
return ToolRegistry([shot, plain])
|
||||
|
||||
|
||||
def _call(name: str) -> AssistantFunctionToolCall:
|
||||
return AssistantFunctionToolCall(id="call-1", function=AssistantFunctionTool(name=name, arguments="{}"))
|
||||
|
||||
|
||||
async def test_execute_keeps_the_text_only_contract() -> None:
|
||||
registry = _registry()
|
||||
|
||||
assert await registry.execute("shot", "{}") == "image: shot.png (2 KB)"
|
||||
assert await registry.execute("plain", "{}") == "just text"
|
||||
|
||||
|
||||
async def test_execute_output_returns_the_images_beside_the_text() -> None:
|
||||
registry = _registry()
|
||||
output = await registry.execute_output("shot", "{}")
|
||||
|
||||
assert output.text == "image: shot.png (2 KB)"
|
||||
assert output.has_images
|
||||
assert output.images == [_image()]
|
||||
|
||||
|
||||
async def test_execute_output_wraps_a_plain_string_result() -> None:
|
||||
output = await _registry().execute_output("plain", "{}")
|
||||
|
||||
assert output.text == "just text"
|
||||
assert not output.has_images
|
||||
|
||||
|
||||
async def test_an_unknown_tool_is_error_text_without_images() -> None:
|
||||
output = await _registry().execute_output("missing", "{}")
|
||||
|
||||
assert "unknown tool" in output.text
|
||||
assert not output.has_images
|
||||
|
||||
|
||||
async def test_the_loop_keeps_tool_images_on_the_tool_message() -> None:
|
||||
message, err = await _run_one_tool(_registry(), _call("shot"), max_result_chars=32_000)
|
||||
|
||||
assert err is None
|
||||
assert isinstance(message, ToolChatMessage)
|
||||
assert message.content == "image: shot.png (2 KB)"
|
||||
assert message.images == [_image()]
|
||||
|
||||
|
||||
async def test_the_loop_leaves_the_images_field_unset_for_text_tools() -> None:
|
||||
message, err = await _run_one_tool(_registry(), _call("plain"), max_result_chars=32_000)
|
||||
|
||||
assert err is None
|
||||
assert message.images is UNSET
|
||||
|
||||
|
||||
async def test_a_truncated_tool_result_keeps_its_images() -> None:
|
||||
message, _ = await _run_one_tool(_registry(), _call("shot"), max_result_chars=1)
|
||||
|
||||
assert "[Truncated" in message.content
|
||||
assert message.images == [_image()]
|
||||
|
||||
|
||||
def test_the_char_estimate_counts_images() -> None:
|
||||
assert estimate_message_chars(UserChatMessage(content="look:")) == len("look:")
|
||||
assert estimate_message_chars(UserChatMessage(content=[ChatTextPart(text="look:"), _image()])) == (
|
||||
len("look:") + IMAGE_CHARS
|
||||
)
|
||||
assert estimate_message_chars(ToolChatMessage(content="shot", tool_call_id="call-1", images=[_image()])) == (
|
||||
len("shot") + len("call-1") + IMAGE_CHARS
|
||||
)
|
||||
Reference in New Issue
Block a user