image filtering bug in _maybe_filter_to_n_most_recent_images?

[rlkiser]

I've been seeing token use growth over time with screenshot-heavy workloads using the computer use demo. I think there's a bug in _maybe_filter_to_n_most_recent_images in loop.py in computer-use-demo that might be part of the problem. I don't exactly understand the caching and chunking that's referenced in the comments, so I'm not comfortable proposing a specific fix here that would preserve intended functionality, but the behavior I'd expect is this:

When a user chooses "Only send N most recent images", the system should remove older images, keeping only the N most recent ones to save on tokens. If there are 5 images and the user specifies 3, it will remove the oldest 2, leaving the most recent 3.

The code to set images_to_remove won't consistently do this however. From line 221:

images_to_remove = total_images - images_to_keep
# for better cache behavior, we want to remove in chunks
images_to_remove -= images_to_remove % min_removal_threshold

Here's how it flows, with line numbers:

103: image_truncation_threshold is set to only_n_most_recent_images or 0
121: When _maybe_filter_to_n_most_recent_images is called, assign only_n_most_recent_images to images_to_keep and image_truncation_threshold to min_removal_threshold in _maybe_filter_to_n_most_recent_images
221: images_to_remove is set to total_images - images_to_keep
223: images_to_remove is set to images_to_remove - (images_to_remove % min_removal_threshold)

For example, assume total images = 5 and only_n_most_recent_images is set to 3 (min_removal_threshold is set to only_n_most_recent_images or 0; lines 125 and 103, default in streamlit.py is 3):

images_to_remove = 5 - 3 = 2
images_to_remove = 2 - 2 % 3 = 0

So the actual behavior in this case would be:
When min_removal_threshold is set to 3, this will set images_to_remove to 0, leaving 5 images where the user specified only 3.

A modulo operation where the numerator is less than the denominator will return the numerator, so whenever images to remove is smaller than the value for only_n_most_recent_images set by streamlit, it will set images_to_remove to 0 regardless of the total count of images. streamlit.py sets the value (3 by default) and I don't see logic there to handle values which will cause this issue, so it doesn't look like this case is handled somewhere else that I'm seeing.

This should be somewhat regularly sending more images than the setting in streamlit, so this seems like a bug, but I don't fully understand the comment at line 222 about this function improving caching performance, so maybe I'm just missing something.

[rlkiser]

Looking at it again this morning with fresh eyes and fresh coffee:

At 112:

    if enable_prompt_caching:
            betas.append(PROMPT_CACHING_BETA_FLAG)
            _inject_prompt_caching(messages)
            # Because cached reads are 10% of the price, we don't think it's
            # ever sensible to break the cache by truncating images
            only_n_most_recent_images = 0
            # Use type ignore to bypass TypedDict check until SDK types are updated
            system["cache_control"] = {"type": "ephemeral"}  # type: ignore

        if only_n_most_recent_images:
            _maybe_filter_to_n_most_recent_images(
                messages,
                only_n_most_recent_images,
                min_removal_threshold=image_truncation_threshold,
            )

This would be bypassing the function in question whenever prompt caching is enabled anyway. Since _maybe_filter_to_n_most_recent_images only gets called there, it looks like _maybe_filter_to_n_most_recent images just won't get called when prompt caching is enabled, I don't see a reason to be concerned about caching performance in the image filtering logic. I think we can do away with line 223 entirely.

I'll find some time to test it and make a PR.

[Luc45]

Also facing this issue. Unbelievable that they ignored such as severe issue that completely kills the Computer Use feature, and even ignored the PR. Crossed fingers they'll circle back to this!