diff --git a/src/plyngent/agent/__init__.py b/src/plyngent/agent/__init__.py index c251724..a115625 100644 --- a/src/plyngent/agent/__init__.py +++ b/src/plyngent/agent/__init__.py @@ -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 diff --git a/src/plyngent/agent/budget.py b/src/plyngent/agent/budget.py index 3693ec0..fc2663b 100644 --- a/src/plyngent/agent/budget.py +++ b/src/plyngent/agent/budget.py @@ -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 diff --git a/src/plyngent/agent/loop.py b/src/plyngent/agent/loop.py index 0147cfd..9a9f819 100644 --- a/src/plyngent/agent/loop.py +++ b/src/plyngent/agent/loop.py @@ -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, diff --git a/src/plyngent/agent/tools.py b/src/plyngent/agent/tools.py index 4cca08f..0ffdb2c 100644 --- a/src/plyngent/agent/tools.py +++ b/src/plyngent/agent/tools.py @@ -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 diff --git a/tests/test_agent/test_tool_output.py b/tests/test_agent/test_tool_output.py new file mode 100644 index 0000000..3788ffa --- /dev/null +++ b/tests/test_agent/test_tool_output.py @@ -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 + )