mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-10-05 20:28:23 +00:00
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:
+66
-33
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user