fix(payment): wire Responses input images into reservation, fix original-detail fallbacks

Address review on PR #680:

- calculate_discounted_max_cost now calls estimate_image_tokens_from_input on
  body["input"], so a Responses request carrying an input_image reserves image
  tokens (previously the estimator was defined but never called, reserving 0).
- _estimate_input_image_tokens now fetches and measures remote original-detail
  images instead of blindly reserving the 36,000-token worst case (~117x
  over-reserve for a 512x512 image), falling back to the worst case only when
  the fetch fails or the file is a file_id reference.
- Broken/undecodable data URLs now reserve the declared detail's worst case
  (36,000 for original) rather than the 85-token low-detail floor: base64 is
  validated and PIL failures fall back explicitly instead of silently using
  _get_image_dimensions' 512x512 default.
- Restore test_discount_cannot_be_dodged_by_hiding_prompt_in_tools (dropped in
  the rebase) and add regression tests for the Responses input-image wiring,
  remote original-detail fetch, and broken-data-URL fallback.
This commit is contained in:
redshift
2026-09-12 12:37:39 +02:00
committed by 9qeklajc
parent 81843e1e24
commit 5920bb69da
2 changed files with 264 additions and 33 deletions
+66 -33
View File
@@ -230,17 +230,24 @@ async def calculate_discounted_max_cost(
# for work the reservation never covered.
prompt_tokens = estimate_prompt_tokens(body)
# Images are billed as tokens by the upstream but carry no text for
# ``estimate_prompt_tokens`` to count, so they are estimated separately and
# added on both the chat (``messages``) and Responses (``input``) paths.
image_tokens = 0
if isinstance(messages, list):
image_tokens = await estimate_image_tokens_in_messages(messages)
if image_tokens > 0:
logger.debug(
"Found images in request",
extra={
"model": model,
"image_tokens": image_tokens,
},
)
prompt_tokens += image_tokens
image_tokens += await estimate_image_tokens_in_messages(messages)
input_data = body.get("input")
if input_data is not None:
image_tokens += await estimate_image_tokens_from_input(input_data)
if image_tokens > 0:
logger.debug(
"Found images in request",
extra={
"model": model,
"image_tokens": image_tokens,
},
)
prompt_tokens += image_tokens
if prompt_tokens > 0:
estimated_prompt_delta_sats = (
@@ -629,34 +636,60 @@ async def estimate_image_tokens_in_messages(messages: list) -> int:
async def _estimate_input_image_tokens(item: dict) -> int:
"""Estimate tokens for a Responses API ``input_image`` item.
Honors the item-level ``detail``. The dimensions of ``file_id``
references can't be fetched here, so they get conservative
estimates: the max-size tile math for high/auto and the 30,000-patch
worst case (36,000 tokens) for original, so we never under-reserve.
Honors the item-level ``detail``. Data-URL images are measured from their
decoded bytes; remote URLs are fetched and measured like the chat path.
Only ``file_id`` references (whose dimensions cannot be fetched here) and
unfetchable/broken images fall back to conservative estimates: the
max-size tile math for high/auto and the 30,000-patch worst case (36,000
tokens) for original, so we never under-reserve.
"""
detail = item.get("detail") or "auto"
if image_url := item.get("image_url"):
if isinstance(image_url, dict):
image_url = image_url.get("url", "")
if isinstance(image_url, str) and image_url.startswith("data:image/"):
try:
_, base64_data = image_url.split(",", 1)
image_bytes = base64.b64decode(base64_data)
width, height = _get_image_dimensions(image_bytes)
return _calculate_image_tokens(width, height, detail)
except Exception as e:
logger.warning(
"Failed to process base64 image", extra={"error": str(e)}
)
return 85
# Remote URLs and file_id both have unfetchable dimensions here; fall
# through to the conservative estimates below.
if item.get("file_id") or item.get("image_url"):
image_url = item.get("image_url")
if isinstance(image_url, dict):
image_url = image_url.get("url", "")
def _worst_case() -> int:
# Dimensions unknown: reserve the worst case for the declared detail so
# a broken/unreadable image still covers what the upstream could bill.
if detail == "original":
return _MAX_ORIGINAL_IMAGE_TOKENS
# We can't fetch an uploaded file's dimensions here; assume the
# largest vision image so we don't under-reserve.
return _calculate_image_tokens(2048, 2048, detail)
if isinstance(image_url, str) and image_url:
if image_url.startswith("data:image/"):
image_bytes = None
try:
_, base64_data = image_url.split(",", 1)
image_bytes = base64.b64decode(base64_data, validate=True)
except Exception as e:
logger.warning(
"Failed to decode base64 image", extra={"error": str(e)}
)
if image_bytes is not None:
try:
img = Image.open(BytesIO(image_bytes))
return _calculate_image_tokens(img.size[0], img.size[1], detail)
except Exception as e:
logger.warning(
"Failed to read image dimensions", extra={"error": str(e)}
)
# Undecodable / unreadable data URL: reserve the worst case.
return _worst_case()
# Remote URL: fetch and measure like the chat path so a small image
# does not reserve the original-detail worst case.
image_bytes = await _fetch_image_from_url(image_url)
if image_bytes:
try:
img = Image.open(BytesIO(image_bytes))
return _calculate_image_tokens(img.size[0], img.size[1], detail)
except Exception as e:
logger.warning(
"Failed to read image dimensions", extra={"error": str(e)}
)
# Unfetchable or unreadable: fall through to the conservative estimate.
if item.get("file_id") or image_url:
return _worst_case()
return 0
+198
View File
@@ -211,6 +211,204 @@ async def test_discount_counts_legacy_token_id_prompt() -> None:
assert cost == 50_000
async def test_discount_cannot_be_dodged_by_hiding_prompt_in_tools() -> None:
"""A large prompt moved from messages into tool schemas must reserve the
same cost — otherwise a caller undercharges by hiding weight from the
estimator."""
from routstr.payment.helpers import calculate_discounted_max_cost
pricing = Mock()
pricing.prompt = 0.5
pricing.completion = 0.01
pricing.max_prompt_cost = 100.0
pricing.max_completion_cost = 100.0
model_obj = Mock()
model_obj.sats_pricing = pricing
model_obj.top_provider = None
model_obj.context_length = None
big_text = "word " * 2_000
base = {"model": "test-model", "max_tokens": 10}
in_messages = {
**base,
"messages": [{"role": "user", "content": big_text}],
}
hiding_places = {
"tools": {
**base,
"messages": [{"role": "user", "content": "hi"}],
"tools": [
{"type": "function", "function": {"name": "f", "description": big_text}}
],
},
# Anthropic forwards a top-level system prompt; it is billed like any other.
"system": {
**base,
"messages": [{"role": "user", "content": "hi"}],
"system": big_text,
},
# A key named like an image field must not win an image exclusion.
"image-named key": {
**base,
"messages": [{"role": "user", "content": "hi"}],
"tools": [{"function": {"parameters": {"data": big_text}}}],
},
# Nor may a caller-chosen "data:" prefix, in any field the body allows.
"data-prefixed content": {
**base,
"messages": [{"role": "user", "content": "data:" + big_text}],
},
"data-prefixed text block": {
**base,
"messages": [
{
"role": "user",
"content": [{"type": "text", "text": "data:" + big_text}],
}
],
},
"data-prefixed system": {
**base,
"messages": [{"role": "user", "content": "hi"}],
"system": "data:" + big_text,
},
}
with (
patch.object(settings, "fixed_pricing", False),
patch.object(settings, "tolerance_percentage", 0),
patch.object(settings, "min_request_msat", 1000),
):
cost_messages = await calculate_discounted_max_cost(
150_000, in_messages, model_obj
)
for where, body in hiding_places.items():
cost = await calculate_discounted_max_cost(150_000, body, model_obj)
# Same prompt weight → at least the same reservation, never the floor.
assert cost >= cost_messages, where
assert cost > 1000, where
async def test_discounted_max_cost_counts_responses_input_images() -> None:
"""A Responses ``input_image`` must add image tokens to the reservation.
Regression for the review finding that ``estimate_image_tokens_from_input``
was defined but never called: a Responses body carries ``input``, not
``messages``, so its images were previously reserved at zero tokens.
"""
import base64
from io import BytesIO
from PIL import Image
from routstr.payment.helpers import calculate_discounted_max_cost
pricing = Mock()
pricing.prompt = 0.001
pricing.completion = 0.001
pricing.max_prompt_cost = 100.0
pricing.max_completion_cost = 0.0
model_obj = Mock()
model_obj.sats_pricing = pricing
model_obj.top_provider = None
model_obj.context_length = None
image = Image.new("RGB", (512, 512), "red")
buffer = BytesIO()
image.save(buffer, format="JPEG")
data_url = "data:image/jpeg;base64," + base64.b64encode(buffer.getvalue()).decode()
no_image = {
"model": "test-model",
"input": [{"role": "user", "content": "hi"}],
}
with_image = {
"model": "test-model",
"input": [
{
"role": "user",
"content": [
{"type": "input_text", "text": "hi"},
{"type": "input_image", "image_url": data_url, "detail": "high"},
],
}
],
}
with (
patch.object(settings, "fixed_pricing", False),
patch.object(settings, "tolerance_percentage", 0),
patch.object(settings, "min_request_msat", 1000),
):
cost_no_image = await calculate_discounted_max_cost(100_000, no_image, model_obj)
cost_with_image = await calculate_discounted_max_cost(
100_000, with_image, model_obj
)
# The 512x512 high-detail image (85 + 170 = 255 tokens) is billed as prompt
# weight, so it reserves strictly more than the identical text-only body.
assert cost_with_image > cost_no_image
async def test_estimate_input_image_tokens_remote_original_fetches() -> None:
"""A remote ``original`` image is fetched and measured, not worst-cased.
Regression for the review finding that any non-data URL with
``detail: \"original\"`` reserved the 36,000-token worst case without
trying to fetch — a 512x512 image reserved ~117x its real cost.
"""
from io import BytesIO
from unittest.mock import patch as mock_patch
from PIL import Image
from routstr.payment.helpers import _estimate_input_image_tokens
image = Image.new("RGB", (512, 512), "red")
buffer = BytesIO()
image.save(buffer, format="JPEG")
image_bytes = buffer.getvalue()
with mock_patch(
"routstr.payment.helpers._fetch_image_from_url",
new=AsyncMock(return_value=image_bytes),
):
# 512x512 original -> 16x16 = 256 patches -> ceil(256 * 1.2) = 308 tokens,
# far below the 36,000 worst case a blind fallback would reserve.
assert await _estimate_input_image_tokens(
{"type": "input_image", "image_url": "https://x.test/i.jpg", "detail": "original"}
) == 308
# When the fetch fails, fall back to the original-detail worst case.
with mock_patch(
"routstr.payment.helpers._fetch_image_from_url",
new=AsyncMock(return_value=None),
):
assert await _estimate_input_image_tokens(
{"type": "input_image", "image_url": "https://x.test/i.jpg", "detail": "original"}
) == 36_000
async def test_estimate_input_image_tokens_broken_original_data_url() -> None:
"""A broken data URL with ``detail: \"original\"`` reserves the worst case.
Regression for the review finding that the ``except`` branch returned 85
(low-detail) regardless of the declared detail.
"""
from routstr.payment.helpers import _estimate_input_image_tokens
# "!!!" is not valid base64, so decoding raises before any dimension read.
assert await _estimate_input_image_tokens(
{"type": "input_image", "image_url": "data:image/jpeg;base64,!!!", "detail": "original"}
) == 36_000
# Non-original details fall back to the max-size tile math, not the 85 floor.
assert await _estimate_input_image_tokens(
{"type": "input_image", "image_url": "data:image/jpeg;base64,!!!", "detail": "high"}
) == 85 + (170 * 4)
async def test_discounted_max_cost_body_max_output_tokens_fallback() -> None:
"""Body ``max_output_tokens`` (Responses API) is honored as a completion cap."""
from routstr.payment.helpers import calculate_discounted_max_cost