fix(payment): honor max_completion_tokens in completion cost reservations

Modern chat-completions clients (newer OpenAI SDKs) send
max_completion_tokens instead of max_tokens; previously such requests got
no completion discount at all and reserved the model's full
max_completion_cost.

Completion caps are now collected from all three spellings (max_tokens,
max_completion_tokens, max_output_tokens) and the reservation is trimmed
by the largest valid one: upstream precedence between the fields varies
by provider, so reserving against a smaller declared cap could
under-cover what the upstream bills. Unparseable values warn and are
skipped, preserving the previous no-discount behavior when no valid cap
is present.
This commit is contained in:
redshift
2026-09-12 12:37:39 +02:00
committed by 9qeklajc
parent c1c3244840
commit 81843e1e24
2 changed files with 141 additions and 71 deletions
+35 -12
View File
@@ -179,7 +179,12 @@ async def calculate_discounted_max_cost(
body: dict,
model_obj: Any | None = None,
) -> int:
"""Calculate the discounted max cost for a request using model pricing when available."""
"""Calculate the discounted max cost for a request using model pricing when available.
Completion discounts are trimmed from the largest declared cap among
``max_tokens`` and ``max_completion_tokens`` (chat/completions) or
``max_output_tokens`` (responses).
"""
if settings.fixed_pricing:
return max_cost_for_model
@@ -244,21 +249,39 @@ async def calculate_discounted_max_cost(
if estimated_prompt_delta_sats > 0:
adjusted = adjusted - math.floor(estimated_prompt_delta_sats * 1000)
max_tokens_raw = body.get("max_tokens", None)
if max_tokens_raw is not None:
# Completion caps arrive under several names: ``max_tokens`` (legacy
# chat), ``max_completion_tokens`` (modern chat) and ``max_output_tokens``
# (Responses API). When a request declares more than one, reserve against
# the largest: upstream precedence between the fields varies by provider,
# so the smaller cap may not be honored and the reservation must never
# under-cover what the upstream could bill.
max_tokens_int: int | None = None
for cap_field in ("max_tokens", "max_completion_tokens", "max_output_tokens"):
cap_raw = body.get(cap_field)
if cap_raw is None:
continue
try:
max_tokens_int = int(max_tokens_raw)
cap_int = int(cap_raw)
except (TypeError, ValueError):
logger.warning(
"Invalid max_tokens; ignoring in cost adjustment",
extra={"max_tokens": str(max_tokens_raw)[:64], "model": model},
"Invalid completion token cap; ignoring in cost adjustment",
extra={
"field": cap_field,
"value": str(cap_raw)[:64],
"model": model,
},
)
else:
estimated_completion_delta_sats = (
max_completion_allowed_sats - max_tokens_int * model_pricing.completion
)
if estimated_completion_delta_sats > 0:
adjusted = adjusted - math.floor(estimated_completion_delta_sats * 1000)
continue
max_tokens_int = (
cap_int if max_tokens_int is None else max(max_tokens_int, cap_int)
)
if max_tokens_int is not None:
estimated_completion_delta_sats = (
max_completion_allowed_sats - max_tokens_int * model_pricing.completion
)
if estimated_completion_delta_sats > 0:
adjusted = adjusted - math.floor(estimated_completion_delta_sats * 1000)
logger.debug(
"Discounted max cost computed",
+106 -59
View File
@@ -211,16 +211,14 @@ 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."""
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
pricing = Mock()
pricing.prompt = 0.5
pricing.completion = 0.01
pricing.max_prompt_cost = 100.0
pricing.prompt = 0.001
pricing.completion = 0.001
pricing.max_prompt_cost = 0.0
pricing.max_completion_cost = 100.0
model_obj = Mock()
@@ -228,51 +226,69 @@ async def test_discount_cannot_be_dodged_by_hiding_prompt_in_tools() -> None:
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,
},
body = {"max_output_tokens": 80_000}
with (
patch.object(settings, "fixed_pricing", False),
patch.object(settings, "tolerance_percentage", 0),
patch.object(settings, "min_request_msat", 1000),
):
cost = await calculate_discounted_max_cost(100_000, body, model_obj)
assert cost == 80_000
async def test_discounted_max_cost_body_max_completion_tokens_fallback() -> None:
"""Body ``max_completion_tokens`` (modern chat) is honored as a completion cap."""
from routstr.payment.helpers import calculate_discounted_max_cost
pricing = Mock()
pricing.prompt = 0.001
pricing.completion = 0.001
pricing.max_prompt_cost = 0.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
body = {"max_completion_tokens": 80_000}
with (
patch.object(settings, "fixed_pricing", False),
patch.object(settings, "tolerance_percentage", 0),
patch.object(settings, "min_request_msat", 1000),
):
cost = await calculate_discounted_max_cost(100_000, body, model_obj)
assert cost == 80_000
async def test_discounted_max_cost_uses_largest_completion_cap() -> None:
"""With several completion caps declared, the largest bounds the reservation.
Upstream precedence between ``max_tokens`` / ``max_completion_tokens`` /
``max_output_tokens`` varies by provider, so reserving against anything
but the largest could under-cover what the upstream bills.
"""
from routstr.payment.helpers import calculate_discounted_max_cost
pricing = Mock()
pricing.prompt = 0.001
pricing.completion = 0.001
pricing.max_prompt_cost = 0.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
body = {
"max_tokens": 50_000,
"max_completion_tokens": 10_000,
"max_output_tokens": 80_000,
}
with (
@@ -280,14 +296,45 @@ async def test_discount_cannot_be_dodged_by_hiding_prompt_in_tools() -> None:
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
cost = await calculate_discounted_max_cost(100_000, body, model_obj)
# 80_000 is the largest declared cap: 100.0 - 80.0 = 20 sats discount.
assert cost == 80_000
async def test_discounted_max_cost_invalid_completion_cap_ignored() -> None:
"""Unparseable caps yield no completion discount rather than under-reserving."""
from routstr.payment.helpers import calculate_discounted_max_cost
pricing = Mock()
pricing.prompt = 0.001
pricing.completion = 0.001
pricing.max_prompt_cost = 0.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
with (
patch.object(settings, "fixed_pricing", False),
patch.object(settings, "tolerance_percentage", 0),
patch.object(settings, "min_request_msat", 1000),
):
# No valid cap at all -> no completion discount.
cost = await calculate_discounted_max_cost(
100_000, {"max_completion_tokens": "sixty-four-k"}, 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
assert cost == 100_000
# An invalid sibling does not poison a valid cap on another field.
cost = await calculate_discounted_max_cost(
100_000,
{"max_tokens": "bad", "max_completion_tokens": 80_000},
model_obj,
)
assert cost == 80_000
async def test_estimate_image_tokens_from_input_detail_and_file_id() -> None: