diff --git a/docs/api/tools/view-image.md b/docs/api/tools/view-image.md index fb6ac42..414a6d8 100644 --- a/docs/api/tools/view-image.md +++ b/docs/api/tools/view-image.md @@ -38,4 +38,29 @@ async with agent.session() as session: 1. The `ViewImageToolProvider` is initialized with an optional `exec_env` parameter 2. If not provided, it automatically detects the execution environment from the session 3. The `view_image` tool reads image files from the execution environment -4. Images are returned as `ImageContentBlock` objects for the model to see +4. Images are resized/compressed to fit configured byte and pixel budgets, then returned as `ImageContentBlock` objects for the model to see +5. The agent keeps only the most recent `max_images_in_context` image blocks in active LLM context; older images are replaced with a short text placeholder + +## Configuration + +Image size limits are set on the code execution provider: + +```python +from stirrup.tools import LocalCodeExecToolProvider, ViewImageToolProvider + +exec_env = LocalCodeExecToolProvider( + max_image_pixels=1_000_000, # default: RESOLUTION_1MP + max_image_bytes=500_000, # default: MAX_IMAGE_BYTES +) +``` + +Bound how many images remain in context on the agent: + +```python +agent = Agent( + client=client, + name="image_viewer", + max_images_in_context=5, # default; pass None to disable + tools=[exec_env, ViewImageToolProvider()], +) +``` diff --git a/src/stirrup/constants.py b/src/stirrup/constants.py index 03a8579..69507ff 100644 --- a/src/stirrup/constants.py +++ b/src/stirrup/constants.py @@ -26,6 +26,11 @@ def __getattr__(name: str) -> Any: # noqa: ANN401 RESOLUTION_1MP = 1_000_000 # 1 megapixel - default max resolution for images RESOLUTION_480P = 640 * 480 # 480p video resolution +# Image context budget (view_image and other ImageContentBlock ingestion) +MAX_IMAGE_BYTES = 500_000 # Max stored bytes per image after resize/compression +DEFAULT_IMAGE_JPEG_QUALITY = 85 # Initial JPEG quality when compressing large images +MAX_IMAGES_IN_CONTEXT = 5 # Max image blocks retained in active LLM context; older ones are replaced with text + # Code execution SANDBOX_TIMEOUT = 60 * 10 # 10 minutes SANDBOX_REQUEST_TIMEOUT = 60 * 3 # 3 minutes diff --git a/src/stirrup/core/agent.py b/src/stirrup/core/agent.py index 71e937c..c96379f 100644 --- a/src/stirrup/core/agent.py +++ b/src/stirrup/core/agent.py @@ -19,6 +19,7 @@ from stirrup.constants import ( AGENT_MAX_TURNS, CONTEXT_SUMMARIZATION_CUTOFF, + MAX_IMAGES_IN_CONTEXT, TURNS_REMAINING_WARNING_THRESHOLD, ) from stirrup.core.cache import CacheManager, CacheState, compute_task_hash @@ -125,6 +126,37 @@ def _handle_text_only_tool_responses(tool_messages: list[ToolMessage]) -> tuple[ return tool_messages, user_messages +_IMAGE_PRUNED_PLACEHOLDER = "[Earlier image removed from context to reduce token usage.]" + + +def _prune_excess_image_blocks(messages: list[ChatMessage], max_images: int | None) -> None: + """Replace oldest image blocks with text when more than ``max_images`` are present. + + Mutates message content in place so repeated ``view_image`` calls do not grow + LLM context without bound. ``None`` disables pruning. + """ + if max_images is None: + return + + image_locations: list[tuple[int, int]] = [] + for msg_idx, message in enumerate(messages): + content = message.content + if not isinstance(content, list): + continue + for block_idx, block in enumerate(content): + if isinstance(block, ImageContentBlock): + image_locations.append((msg_idx, block_idx)) + + excess = len(image_locations) - max_images + if excess <= 0: + return + + for msg_idx, block_idx in image_locations[:excess]: + message = messages[msg_idx] + if isinstance(message.content, list): + message.content[block_idx] = _IMAGE_PRUNED_PLACEHOLDER + + def _normalize_finish_tools[FinishParams: BaseModel, FinishMeta]( finish_tool: Tool[FinishParams, FinishMeta] | list[Tool] | None, ) -> dict[str, Tool[Any, Any]]: @@ -277,6 +309,7 @@ def __init__( text_only_tool_responses: bool = ..., block_successive_assistant_messages: bool = ..., recover_from_context_overflow: bool = ..., + max_images_in_context: int | None = ..., share_parent_exec_env: bool = ..., logger: AgentLoggerBase | None = ..., ) -> None: ... @@ -297,6 +330,7 @@ def __init__( text_only_tool_responses: bool = ..., block_successive_assistant_messages: bool = ..., recover_from_context_overflow: bool = ..., + max_images_in_context: int | None = ..., share_parent_exec_env: bool = ..., logger: AgentLoggerBase | None = ..., ) -> None: ... @@ -317,6 +351,7 @@ def __init__( text_only_tool_responses: bool = ..., block_successive_assistant_messages: bool = ..., recover_from_context_overflow: bool = ..., + max_images_in_context: int | None = ..., share_parent_exec_env: bool = ..., logger: AgentLoggerBase | None = ..., ) -> None: ... @@ -337,6 +372,7 @@ def __init__( text_only_tool_responses: bool = True, block_successive_assistant_messages: bool = True, recover_from_context_overflow: bool = True, + max_images_in_context: int | None = MAX_IMAGES_IN_CONTEXT, # Subagent options share_parent_exec_env: bool = False, # Logging @@ -366,6 +402,9 @@ def __init__( recover_from_context_overflow: If True (default), drop one completed turn and retry when the model reports a context overflow. If the original prompt still overflows, the context error is raised. + max_images_in_context: Maximum number of image blocks to keep in active LLM context. + Older images are replaced with a short text placeholder before + each model call. ``None`` disables this limit. share_parent_exec_env: When True and used as a subagent, share the parent's code execution environment instead of creating a new one. This provides better performance (no file copying) and allows @@ -382,6 +421,9 @@ def __init__( "(alphanumeric, underscores, hyphens only, 1-128 characters)." ) + if max_images_in_context is not None and max_images_in_context < 0: + raise ValueError("max_images_in_context must be non-negative or None") + self._client: LLMClient = client self._name = name self._max_turns = max_turns @@ -394,6 +436,7 @@ def __init__( self._text_only_tool_responses = text_only_tool_responses self._block_successive_assistant_messages = block_successive_assistant_messages self._recover_from_context_overflow = recover_from_context_overflow + self._max_images_in_context = max_images_in_context self._share_parent_exec_env = share_parent_exec_env # Logger (can be passed in or created here) @@ -1216,6 +1259,7 @@ async def step( Returns the assistant message, tool execution results, and finish tool call (if present). """ + _prune_excess_image_blocks(messages, self._max_images_in_context) assistant_message = await self._client.generate(messages, self._active_tools) # Log assistant message immediately @@ -1279,6 +1323,7 @@ async def summarize_messages( ) summary_prompt = [*current_messages, UserMessage(content=MESSAGE_SUMMARIZER)] + _prune_excess_image_blocks(summary_prompt, self._max_images_in_context) # Give the summarizer the active tools so it can interpret prior tool calls/results. summary = await self._client.generate(summary_prompt, self._active_tools) diff --git a/src/stirrup/core/models.py b/src/stirrup/core/models.py index 8b6eb44..a9abbf5 100644 --- a/src/stirrup/core/models.py +++ b/src/stirrup/core/models.py @@ -19,7 +19,12 @@ from PIL import Image from pydantic import BaseModel, Field, PlainSerializer, PlainValidator, model_validator -from stirrup.constants import RESOLUTION_1MP, RESOLUTION_480P +from stirrup.constants import ( + DEFAULT_IMAGE_JPEG_QUALITY, + MAX_IMAGE_BYTES, + RESOLUTION_1MP, + RESOLUTION_480P, +) __all__ = [ "Addable", @@ -46,6 +51,7 @@ "UserMessage", "VideoContentBlock", "aggregate_metadata", + "prepare_image_bytes", ] @@ -78,6 +84,77 @@ def downscale_image(w: int, h: int, max_pixels: int | None = 1_000_000) -> tuple return max(nw, 2), max(nh, 2) +def prepare_image_bytes( + data: bytes, + *, + max_pixels: int | None = RESOLUTION_1MP, + max_bytes: int | None = MAX_IMAGE_BYTES, + jpeg_quality: int = DEFAULT_IMAGE_JPEG_QUALITY, +) -> bytes: + """Resize and compress image bytes before storing them in agent context. + + Images that already fit within ``max_pixels`` and ``max_bytes`` are returned unchanged. + Larger images are downscaled to ``max_pixels`` (when set) and re-encoded as JPEG. + When quality reduction alone cannot satisfy ``max_bytes``, dimensions are reduced + further (down to a 2x2 lower bound) until the byte budget is met or no further + reduction is possible. + + Args: + data: Raw image file bytes. + max_pixels: Maximum pixel count (width * height). ``None`` disables pixel limits. + max_bytes: Maximum encoded byte size. ``None`` disables byte limits. + jpeg_quality: Starting JPEG quality used when re-encoding oversized images. + + Returns: + Image bytes suitable for ``ImageContentBlock`` storage and LLM context. + """ + min_dimension = 2 + + with Image.open(BytesIO(data)) as img: + width, height = img.size + needs_resize = max_pixels is not None and width * height > max_pixels + needs_compress = max_bytes is not None and len(data) > max_bytes + if not needs_resize and not needs_compress: + return data + + working = img.copy() + if needs_resize and max_pixels is not None: + target_w, target_h = downscale_image(width, height, max_pixels) + working.thumbnail((target_w, target_h), Image.Resampling.LANCZOS) + + if working.mode != "RGB": + working = working.convert("RGB") + + if max_bytes is None: + buf = BytesIO() + working.save(buf, format="JPEG", quality=jpeg_quality, optimize=True) + return buf.getvalue() + + quality = jpeg_quality + while True: + buf = BytesIO() + working.save(buf, format="JPEG", quality=quality, optimize=True) + result = buf.getvalue() + if len(result) <= max_bytes: + return result + + if quality > 30: + quality -= 10 + continue + + current_w, current_h = working.size + if current_w <= min_dimension or current_h <= min_dimension: + return result + + new_w = max(min_dimension, current_w // 2) + new_h = max(min_dimension, current_h // 2) + if (new_w, new_h) == (current_w, current_h): + return result + + working = working.resize((new_w, new_h), Image.Resampling.LANCZOS) + quality = jpeg_quality + + # Content class BinaryContentBlock(BaseModel, ABC): """Base class for binary content (images, video, audio) with MIME type validation.""" @@ -115,7 +192,11 @@ def _probe(self) -> None: class ImageContentBlock(BinaryContentBlock): - """Image content supporting PNG, JPEG, WebP, PSD formats with automatic downscaling.""" + """Image content supporting PNG, JPEG, WebP, PSD formats with automatic downscaling. + + Use :func:`prepare_image_bytes` before constructing blocks from large files (for example + via ``view_image``) so stored bytes stay within configured pixel and byte budgets. + """ kind: Literal["image_content_block"] = "image_content_block" allowed_mime_types: ClassVar[set[str]] = { diff --git a/src/stirrup/tools/code_backends/base.py b/src/stirrup/tools/code_backends/base.py index 73fa015..21b8009 100644 --- a/src/stirrup/tools/code_backends/base.py +++ b/src/stirrup/tools/code_backends/base.py @@ -9,7 +9,15 @@ from pydantic import BaseModel, Field -from stirrup.core.models import ImageContentBlock, Tool, ToolProvider, ToolResult, ToolUseCountMetadata +from stirrup.constants import MAX_IMAGE_BYTES, RESOLUTION_1MP +from stirrup.core.models import ( + ImageContentBlock, + Tool, + ToolProvider, + ToolResult, + ToolUseCountMetadata, + prepare_image_bytes, +) from stirrup.utils.text import truncate_msg logger = logging.getLogger(__name__) @@ -151,6 +159,8 @@ def __init__( *, allowed_commands: list[str] | None = None, shell_timeout: int = SHELL_TIMEOUT, + max_image_pixels: int | None = RESOLUTION_1MP, + max_image_bytes: int | None = MAX_IMAGE_BYTES, ) -> None: """Initialize execution environment with optional command allowlist. @@ -163,10 +173,16 @@ def __init__( to ``SHELL_TIMEOUT``. Callers should set this to match their application's expected long-running-command budget rather than rely on the stirrup default. + max_image_pixels: Maximum pixel count (width * height) for images returned + by ``view_image``. ``None`` disables pixel limits. + max_image_bytes: Maximum encoded byte size for images returned by + ``view_image``. ``None`` disables byte limits. """ self._allowed_commands = allowed_commands self._shell_timeout = shell_timeout + self._max_image_pixels = max_image_pixels + self._max_image_bytes = max_image_bytes self._compiled_allowed: list[re.Pattern[str]] | None = None if allowed_commands is not None: self._compiled_allowed = [re.compile(p) for p in allowed_commands] @@ -508,10 +524,18 @@ def get_view_image_tool( """ env = self + max_image_pixels = self._max_image_pixels + max_image_bytes = self._max_image_bytes async def executor(params: ViewImageParams) -> ToolResult[ToolUseCountMetadata]: try: image = await env.view_image(params.path) + normalized_data = prepare_image_bytes( + image.data, + max_pixels=max_image_pixels, + max_bytes=max_image_bytes, + ) + image = ImageContentBlock(data=normalized_data) return ToolResult( content=["Viewing image at path: " + params.path, image], metadata=ToolUseCountMetadata(), diff --git a/src/stirrup/tools/code_backends/local.py b/src/stirrup/tools/code_backends/local.py index 247178e..7811142 100644 --- a/src/stirrup/tools/code_backends/local.py +++ b/src/stirrup/tools/code_backends/local.py @@ -12,6 +12,7 @@ import anyio +from stirrup.constants import MAX_IMAGE_BYTES, RESOLUTION_1MP from stirrup.core.models import ImageContentBlock, Tool, ToolUseCountMetadata from .base import ( @@ -63,6 +64,8 @@ def __init__( temp_base_dir: Path | str | None = None, description: str | None = None, shell_timeout: int = SHELL_TIMEOUT, + max_image_pixels: int | None = RESOLUTION_1MP, + max_image_bytes: int | None = MAX_IMAGE_BYTES, ) -> None: """Initialize LocalCodeExecToolProvider configuration. @@ -77,9 +80,16 @@ def __init__( ``code_exec`` invocation, and the default for direct ``run_command`` calls that pass ``timeout=None``. Defaults to ``SHELL_TIMEOUT``. + max_image_pixels: Maximum pixel count for ``view_image`` results. + max_image_bytes: Maximum encoded byte size for ``view_image`` results. """ - super().__init__(allowed_commands=allowed_commands, shell_timeout=shell_timeout) + super().__init__( + allowed_commands=allowed_commands, + shell_timeout=shell_timeout, + max_image_pixels=max_image_pixels, + max_image_bytes=max_image_bytes, + ) self._temp_dir: Path | None = None self._temp_base_dir: Path | None = Path(temp_base_dir) if temp_base_dir else None self._description = ( diff --git a/src/stirrup/tools/view_image.py b/src/stirrup/tools/view_image.py index eb20282..efa4c4f 100644 --- a/src/stirrup/tools/view_image.py +++ b/src/stirrup/tools/view_image.py @@ -7,6 +7,15 @@ class ViewImageToolProvider(ToolProvider): """Tool provider for viewing images from an execution environment. + Images returned by ``view_image`` are resized and compressed according to the + execution environment's ``max_image_pixels`` and ``max_image_bytes`` settings + (defaults: :data:`~stirrup.constants.RESOLUTION_1MP` and + :data:`~stirrup.constants.MAX_IMAGE_BYTES`). Configure these on the + ``CodeExecToolProvider`` instance (for example ``LocalCodeExecToolProvider``). + + The agent also bounds how many image blocks stay in active LLM context via + ``max_images_in_context`` (default: :data:`~stirrup.constants.MAX_IMAGES_IN_CONTEXT`). + Can be used with an explicit exec_env or will auto-detect from the Agent's session state. Works regardless of tool ordering in the Agent. diff --git a/tests/test_image_context.py b/tests/test_image_context.py new file mode 100644 index 0000000..7d33b0d --- /dev/null +++ b/tests/test_image_context.py @@ -0,0 +1,242 @@ +"""Tests for image context budgeting (resize/compression and bounded history).""" + +import inspect +import random +from io import BytesIO + +import pytest +from PIL import Image + +from stirrup.constants import MAX_IMAGE_BYTES, RESOLUTION_1MP +from stirrup.core.agent import _prune_excess_image_blocks +from stirrup.core.models import ImageContentBlock, ToolMessage, UserMessage, prepare_image_bytes +from stirrup.tools.code_backends.base import ViewImageParams +from stirrup.tools.code_backends.local import LocalCodeExecToolProvider + + +def _large_image_bytes(*, width: int = 2000, height: int = 2000) -> bytes: + """Create a large uncompressed BMP image for budget tests.""" + img = Image.new("RGB", (width, height), color=(12, 34, 56)) + buffer = BytesIO() + img.save(buffer, format="BMP") + return buffer.getvalue() + + +def _noisy_image_bytes(*, width: int = 256, height: int = 256) -> bytes: + """Create a high-entropy image that resists quality-only JPEG compression.""" + img = Image.new("RGB", (width, height)) + rng = random.Random(0) + for y in range(height): + for x in range(width): + img.putpixel((x, y), (rng.randint(0, 255), rng.randint(0, 255), rng.randint(0, 255))) + buffer = BytesIO() + img.save(buffer, format="PNG") + return buffer.getvalue() + + +class TestPrepareImageBytes: + """Tests for prepare_image_bytes normalization.""" + + def test_preserves_small_images(self, sample_png_bytes: bytes) -> None: + """Small images should pass through unchanged.""" + assert prepare_image_bytes(sample_png_bytes) == sample_png_bytes + + def test_downscales_and_compresses_large_images(self) -> None: + """Large images should be reduced below configured byte and pixel limits.""" + raw = _large_image_bytes() + assert len(raw) > MAX_IMAGE_BYTES + + prepared = prepare_image_bytes(raw, max_pixels=RESOLUTION_1MP, max_bytes=MAX_IMAGE_BYTES) + assert len(prepared) <= MAX_IMAGE_BYTES + + with Image.open(BytesIO(prepared)) as img: + assert img.width * img.height <= RESOLUTION_1MP + + def test_respects_disabled_limits(self) -> None: + """Passing None for limits should preserve original bytes.""" + raw = _large_image_bytes(width=10, height=10) + assert prepare_image_bytes(raw, max_pixels=None, max_bytes=None) == raw + + def test_noisy_image_meets_byte_budget_via_dimension_reduction(self) -> None: + """Noisy images should still satisfy max_bytes after dimension reduction.""" + raw = _noisy_image_bytes() + max_bytes = 8_000 + + buf = BytesIO() + with Image.open(BytesIO(raw)) as img: + rgb = img.convert("RGB") + rgb.save(buf, format="JPEG", quality=30, optimize=True) + assert len(buf.getvalue()) > max_bytes + + prepared = prepare_image_bytes(raw, max_pixels=None, max_bytes=max_bytes) + assert len(prepared) <= max_bytes + + +class TestViewImageNormalization: + """Tests for view_image tool ingestion normalization.""" + + async def test_view_image_tool_compresses_large_file(self) -> None: + """view_image should store normalized bytes, not the full source file.""" + provider = LocalCodeExecToolProvider(max_image_bytes=50_000) + + async with provider: + assert provider._temp_dir is not None # noqa: SLF001 + image_path = provider._temp_dir / "large.png" # noqa: SLF001 + image_path.write_bytes(_large_image_bytes()) + + tool = provider.get_view_image_tool() + params = ViewImageParams(path="large.png") + executor_result = tool.executor(params) + result = await executor_result if inspect.isawaitable(executor_result) else executor_result + + assert isinstance(result.content, list) + image_block = result.content[1] + assert isinstance(image_block, ImageContentBlock) + assert len(image_block.data) <= 50_000 + + async def test_repeated_view_image_calls_stay_bounded(self) -> None: + """Repeated view_image calls should each return bounded image payloads.""" + provider = LocalCodeExecToolProvider(max_image_bytes=40_000) + + async with provider: + assert provider._temp_dir is not None # noqa: SLF001 + image_path = provider._temp_dir / "chart.png" # noqa: SLF001 + image_path.write_bytes(_large_image_bytes()) + + tool = provider.get_view_image_tool() + params = ViewImageParams(path="chart.png") + + for _ in range(3): + executor_result = tool.executor(params) + result = await executor_result if inspect.isawaitable(executor_result) else executor_result + image_block = result.content[1] + assert isinstance(image_block, ImageContentBlock) + assert len(image_block.data) <= 40_000 + + +class TestImageHistoryPruning: + """Tests for bounded image history before LLM calls.""" + + async def test_agent_step_prunes_images_before_generate(self) -> None: + """Agent.step should not send every historical image block to the model.""" + from stirrup import Agent + from stirrup.core.models import AssistantMessage, ChatMessage, TokenUsage + from tests.test_agent import MockLLMClient + + seen_image_blocks = 0 + + class CountingClient(MockLLMClient): + async def generate( + self, + messages: list[ChatMessage], + tools: object, + ) -> AssistantMessage: + del tools + nonlocal seen_image_blocks + for message in messages: + if isinstance(message.content, list): + seen_image_blocks += sum(1 for block in message.content if isinstance(block, ImageContentBlock)) + return AssistantMessage( + content="ok", + tool_calls=[], + token_usage=TokenUsage(input=1, answer=1), + ) + + blocks = [ImageContentBlock(data=_large_image_bytes(width=2, height=2)) for _ in range(4)] + messages = [ + ToolMessage(content=["older", blocks[0]], tool_call_id="call_0", name="view_image"), + ToolMessage(content=["older", blocks[1]], tool_call_id="call_1", name="view_image"), + ToolMessage(content=["newer", blocks[2]], tool_call_id="call_2", name="view_image"), + ToolMessage(content=["newest", blocks[3]], tool_call_id="call_3", name="view_image"), + ] + + agent = Agent( + client=CountingClient([]), + name="test-agent", + max_turns=1, + max_images_in_context=2, + tools=[], + ) + + await agent.step(messages, run_metadata={}) + assert seen_image_blocks == 2 + + def test_prune_replaces_oldest_images(self) -> None: + """Only the newest max_images blocks should remain as images.""" + blocks = [ImageContentBlock(data=_large_image_bytes(width=2, height=2)) for _ in range(4)] + messages = [ + ToolMessage( + content=["first", blocks[0]], + tool_call_id="call_1", + name="view_image", + ), + UserMessage(content=[blocks[1]]), + ToolMessage( + content=["second", blocks[2]], + tool_call_id="call_2", + name="view_image", + ), + UserMessage(content=[blocks[3]]), + ] + + _prune_excess_image_blocks(messages, max_images=2) + + remaining_images = [ + block + for message in messages + if isinstance(message.content, list) + for block in message.content + if isinstance(block, ImageContentBlock) + ] + assert len(remaining_images) == 2 + + first_tool = messages[0] + assert isinstance(first_tool.content, list) + assert isinstance(first_tool.content[1], str) + assert "removed from context" in first_tool.content[1] + + def test_prune_noop_when_disabled(self) -> None: + """None max_images should leave all image blocks intact.""" + block = ImageContentBlock(data=_large_image_bytes(width=2, height=2)) + messages = [UserMessage(content=[block])] + _prune_excess_image_blocks(messages, max_images=None) + content = messages[0].content + assert isinstance(content, list) + assert isinstance(content[0], ImageContentBlock) + + +class TestAgentImageContextConfig: + """Tests for agent-level image context configuration.""" + + def test_rejects_negative_max_images_in_context(self) -> None: + """max_images_in_context must be non-negative when set.""" + from stirrup import Agent + from tests.test_agent import MockLLMClient + + with pytest.raises(ValueError, match="max_images_in_context must be non-negative or None"): + Agent( + client=MockLLMClient([]), + name="test-agent", + max_images_in_context=-1, + ) + + def test_allows_none_max_images_in_context(self) -> None: + """None should disable image-count pruning.""" + from stirrup import Agent + from tests.test_agent import MockLLMClient + + agent = Agent( + client=MockLLMClient([]), + name="test-agent", + max_images_in_context=None, + ) + assert agent._max_images_in_context is None # noqa: SLF001 + + +@pytest.fixture +def sample_png_bytes() -> bytes: + """Create valid PNG image bytes using PIL (1x1 red pixel).""" + img = Image.new("RGB", (1, 1), color=(255, 0, 0)) + buffer = BytesIO() + img.save(buffer, format="PNG") + return buffer.getvalue()