fix(decisions): report provider charges on the cost header and read Vercel's cache split from the serving provider

Backport of #50642 to lts/1.108.x
Cherry-picked from ec26925103 (mishdotdev:vercel-ai-gateway-decisions)

The Jev 1.13.0 and 1.14.0 rows record their image probes

Conflicts resolved to main's result. The call.py hunk is dropped because the line already reports provider charges on the response-cost header. The proxy test keeps only its paid call, since the line's proxy drops a zero response-cost header for every provider
This commit is contained in:
mateo-berri 2026-11-09 19:26:13 -08:00
parent 9bed368bfd
commit 9f4a342c8b
5 changed files with 146 additions and 44 deletions

View file

@ -1,13 +1,18 @@
from collections.abc import Mapping
from dataclasses import replace
from itertools import chain
from typing import Final, Literal, TypeVar
from typing import Final, TypeVar
from pydantic import BaseModel, ConfigDict, Field, ValidationError
from litellm.llms.base_llm.decisions.transformation import usage_reported_cost
from litellm.llms.openai.decisions.transformation import OpenAIDecisionsConfig
from litellm.types.decisions import DecisionsIRRequest, DecisionsIRResponse, DecisionsIRUsage, OpenAIDecisionUsage
from litellm.types.decisions import (
DecisionsIRRequest,
DecisionsIRResponse,
DecisionsIRUsage,
OpenAIDecisionInputTokensDetails,
)
_Model: Final = TypeVar("_Model", bound=BaseModel)
@ -15,7 +20,8 @@ _Model: Final = TypeVar("_Model", bound=BaseModel)
class _ProviderAttempt(BaseModel):
model_config = ConfigDict(frozen=True)
credential_type: Literal["system", "byok"] | None = Field(default=None, alias="credentialType")
provider: str | None = None
credential_type: str | None = Field(default=None, alias="credentialType")
success: bool = False
@ -30,28 +36,40 @@ class _Routing(BaseModel):
model_attempts: tuple[_ModelAttempt, ...] = Field(default=(), alias="modelAttempts")
def served_by(self) -> _ProviderAttempt | None:
attempts: Final = chain.from_iterable(model.provider_attempts for model in self.model_attempts)
return next((attempt for attempt in attempts if attempt.success), None)
class _GatewayRouting(BaseModel):
model_config = ConfigDict(frozen=True)
routing: _Routing = _Routing()
class _GatewayCharge(BaseModel):
model_config = ConfigDict(frozen=True)
cost: float
market_cost: float | None = Field(default=None, alias="marketCost")
routing: _Routing = _Routing()
def served_with_own_key(self) -> bool:
attempts: Final = chain.from_iterable(model.provider_attempts for model in self.routing.model_attempts)
return any(attempt.success and attempt.credential_type == "byok" for attempt in attempts)
def billed_cost(self) -> float:
if self.market_cost is None or not self.served_with_own_key():
def billed_cost(self, served_by: _ProviderAttempt | None) -> float:
if self.market_cost is None or served_by is None or served_by.credential_type != "byok":
return self.cost
return self.cost + self.market_cost
class _UpstreamUsage(BaseModel):
model_config = ConfigDict(frozen=True)
input_tokens: int
input_tokens_details: OpenAIDecisionInputTokensDetails
class _UpstreamReport(BaseModel):
model_config = ConfigDict(frozen=True)
usage: OpenAIDecisionUsage
usage: _UpstreamUsage
class _GatewayDecisionResponse(BaseModel):
@ -67,15 +85,27 @@ def _validated(model: type[_Model], value: object) -> _Model | None:
return None
def _served_by(gateway: object) -> _ProviderAttempt | None:
gateway_routing: Final = _validated(_GatewayRouting, gateway)
return None if gateway_routing is None else gateway_routing.routing.served_by()
def _upstream_usage(metadata: Mapping[str, object], served_by: _ProviderAttempt | None) -> _UpstreamUsage | None:
if served_by is None or served_by.provider is None:
return None
report: Final = _validated(_UpstreamReport, metadata.get(served_by.provider))
return None if report is None else report.usage
def _gateway_usage(usage: DecisionsIRUsage, metadata: Mapping[str, object]) -> DecisionsIRUsage:
charge: Final = _validated(_GatewayCharge, metadata.get("gateway"))
reports: Final = (_validated(_UpstreamReport, entry) for entry in metadata.values())
upstream: Final = next(
(report.usage for report in reports if report is not None and report.usage.input_tokens == usage.input_tokens),
None,
gateway: Final = metadata.get("gateway")
served_by: Final = _served_by(gateway)
charge: Final = _validated(_GatewayCharge, gateway)
priced: Final = (
usage if charge is None else replace(usage, extra={**usage.extra, "cost": charge.billed_cost(served_by)})
)
priced: Final = usage if charge is None else replace(usage, extra={**usage.extra, "cost": charge.billed_cost()})
if upstream is None:
upstream: Final = _upstream_usage(metadata, served_by)
if upstream is None or upstream.input_tokens != usage.input_tokens:
return priced
return replace(
priced,

View file

@ -51286,7 +51286,7 @@
"litellm_provider": "vercel_ai_gateway",
"max_input_tokens": 32000,
"metadata": {
"notes": "Vercel AI Gateway model list entry typesafe-ai/jev-1.13.0, read 2026-11-09: input 0.000000042 per token, output 0, context window 32000."
"notes": "Vercel AI Gateway model list entry typesafe-ai/jev-1.13.0, read 2026-11-09: input 0.000000042 per token, output 0, context window 32000. Direct /v1/decisions probes on 2026-11-10 asking whether a 768x768 picture is a smooth gradient answered 0.99 for a gradient and 0.12 for solid gray."
},
"mode": "evaluation",
"output_cost_per_token": 0.0,
@ -51302,7 +51302,7 @@
"litellm_provider": "vercel_ai_gateway",
"max_input_tokens": 32000,
"metadata": {
"notes": "Vercel AI Gateway model list entry typesafe-ai/jev-1.14.0, read 2026-11-09: input 0.000000042 per token, output 0, context window 32000."
"notes": "Vercel AI Gateway model list entry typesafe-ai/jev-1.14.0, read 2026-11-09: input 0.000000042 per token, output 0, context window 32000. Direct /v1/decisions probes on 2026-11-10 asking whether a 768x768 picture is a smooth gradient answered 0.84 for a gradient and 0.13 for solid gray."
},
"mode": "evaluation",
"output_cost_per_token": 0.0,

View file

@ -51286,7 +51286,7 @@
"litellm_provider": "vercel_ai_gateway",
"max_input_tokens": 32000,
"metadata": {
"notes": "Vercel AI Gateway model list entry typesafe-ai/jev-1.13.0, read 2026-11-09: input 0.000000042 per token, output 0, context window 32000."
"notes": "Vercel AI Gateway model list entry typesafe-ai/jev-1.13.0, read 2026-11-09: input 0.000000042 per token, output 0, context window 32000. Direct /v1/decisions probes on 2026-11-10 asking whether a 768x768 picture is a smooth gradient answered 0.99 for a gradient and 0.12 for solid gray."
},
"mode": "evaluation",
"output_cost_per_token": 0.0,
@ -51302,7 +51302,7 @@
"litellm_provider": "vercel_ai_gateway",
"max_input_tokens": 32000,
"metadata": {
"notes": "Vercel AI Gateway model list entry typesafe-ai/jev-1.14.0, read 2026-11-09: input 0.000000042 per token, output 0, context window 32000."
"notes": "Vercel AI Gateway model list entry typesafe-ai/jev-1.14.0, read 2026-11-09: input 0.000000042 per token, output 0, context window 32000. Direct /v1/decisions probes on 2026-11-10 asking whether a 768x768 picture is a smooth gradient answered 0.84 for a gradient and 0.13 for solid gray."
},
"mode": "evaluation",
"output_cost_per_token": 0.0,

View file

@ -3575,36 +3575,65 @@ async def test_vercel_ai_gateway_jev_sends_box_questions(
assert response.answers[0].confidence == 0.8
_UPSTREAM_SPLIT: Final[Mapping[str, int]] = {"cached_tokens": _INPUT_TOKENS - 20, "cache_write_tokens": 20}
_TOP_LEVEL_SPLIT: Final[Mapping[str, int]] = {"cached_tokens": 300, "cache_write_tokens": 7}
def _upstream_usage(input_tokens: int, details: Mapping[str, int] | None) -> Mapping[str, object]:
return {
"usage": {
"input_tokens": input_tokens,
**({} if details is None else {"input_tokens_details": details}),
"output_tokens": 0,
"total_tokens": input_tokens,
}
}
@pytest.mark.asyncio
@pytest.mark.usefixtures("vercel_gateway_key")
@pytest.mark.parametrize(
("upstream_input_tokens", "cached_tokens", "cache_write_tokens"),
((_INPUT_TOKENS, _INPUT_TOKENS - 20, 20), (_INPUT_TOKENS + 1, 0, 0)),
ids=("upstream_reports_this_call", "upstream_count_disagrees"),
("top_level_split", "served_by", "upstream", "expected_split"),
(
(None, "openai", _upstream_usage(_INPUT_TOKENS, _UPSTREAM_SPLIT), _UPSTREAM_SPLIT),
(None, "openai", _upstream_usage(_INPUT_TOKENS + 1, _UPSTREAM_SPLIT), {"cached_tokens": 0, "cache_write_tokens": 0}),
(_TOP_LEVEL_SPLIT, "openai", _upstream_usage(_INPUT_TOKENS, None), _TOP_LEVEL_SPLIT),
(None, "anthropic", _upstream_usage(_INPUT_TOKENS, _UPSTREAM_SPLIT), {"cached_tokens": 0, "cache_write_tokens": 0}),
),
ids=(
"serving_provider_reports_this_call",
"upstream_count_disagrees",
"upstream_omits_the_split",
"report_from_a_provider_that_did_not_serve",
),
)
async def test_vercel_ai_gateway_decisions_report_the_upstream_cache_split(
upstream_input_tokens: int,
cached_tokens: int,
cache_write_tokens: int,
monkeypatch: pytest.MonkeyPatch,
top_level_split: Mapping[str, int] | None,
served_by: str,
upstream: Mapping[str, object],
expected_split: Mapping[str, int],
respx_mock: respx.MockRouter,
) -> None:
monkeypatch.delenv("VERCEL_AI_GATEWAY_API_BASE", raising=False)
monkeypatch.setenv("VERCEL_AI_GATEWAY_API_KEY", "gateway-key")
respx_mock.post("https://ai-gateway.vercel.sh/v1/decisions").respond(
respx_mock.post(_VERCEL_DECISIONS_URL).respond(
json={
"model": "openai/gpt-6-luna-decisions",
"answers": [{"type": "predicate", "name": "is_defect", "probability": 0.9}],
"usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0, "total_tokens": _INPUT_TOKENS},
"usage": {
"input_tokens": _INPUT_TOKENS,
**({} if top_level_split is None else {"input_tokens_details": top_level_split}),
"output_tokens": 0,
"total_tokens": _INPUT_TOKENS,
},
"provider_metadata": {
"openai": {
"usage": {
"input_tokens": upstream_input_tokens,
"input_tokens_details": {"cached_tokens": _INPUT_TOKENS - 20, "cache_write_tokens": 20},
"output_tokens": 0,
"total_tokens": upstream_input_tokens,
}
"openai": upstream,
"gateway": {
"cost": "0.00005722",
"routing": {
"modelAttempts": [
{"providerAttempts": [{"provider": served_by, "credentialType": "system", "success": True}]}
]
},
},
"gateway": {"cost": "0.00005722"},
},
}
)
@ -3612,9 +3641,9 @@ async def test_vercel_ai_gateway_decisions_report_the_upstream_cache_split(
response: Final = await litellm.adecisions(
model="vercel_ai_gateway/openai/gpt-6-luna-decisions",
input="The package arrived broken.",
questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}],
questions=_PREDICATE_QUESTIONS,
)
assert response.usage.input_tokens_details.cached_tokens == cached_tokens
assert response.usage.input_tokens_details.cache_write_tokens == cache_write_tokens
assert response.usage.input_tokens_details.cached_tokens == expected_split["cached_tokens"]
assert response.usage.input_tokens_details.cache_write_tokens == expected_split["cache_write_tokens"]
assert get_response_cost_from_hidden_params(response.hidden_params) == pytest.approx(5.722e-05)

View file

@ -1338,3 +1338,46 @@ def test_proxy_decisions_dispatches_vercel_ai_gateway_deployment(
)
assert expected_cost > 0
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected_cost)
@pytest.mark.parametrize("gateway_cost", ("0.00005722",), ids=("paid_call",))
def test_proxy_decisions_report_the_billed_gateway_charge_in_the_header_and_body(
gateway_cost: str,
client: TestClient,
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
monkeypatch.delenv("VERCEL_AI_GATEWAY_API_BASE", raising=False)
monkeypatch.setenv("VERCEL_AI_GATEWAY_API_KEY", "gateway-key")
monkeypatch.setattr(litellm, "cost_margin_config", {"vercel_ai_gateway": 0.5})
deployment_router: Final = litellm.Router(
model_list=[
{
"model_name": "gateway-decider",
"litellm_params": {"model": "vercel_ai_gateway/example/unpriced-decider"},
"model_info": {"mode": "evaluation"},
}
]
)
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", deployment_router)
respx_mock.post("https://ai-gateway.vercel.sh/v1/decisions").respond(
json={
"model": "example/unpriced-decider",
"answers": [{"type": "predicate", "name": "is_defect", "probability": 0.9}],
"usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0, "total_tokens": _INPUT_TOKENS},
"provider_metadata": {"gateway": {"cost": gateway_cost}},
}
)
response: Final = client.post(
"/v1/systemone",
json={
"model": "gateway-decider",
"state": {"source": "proxy-test"},
"questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}},
},
)
assert response.status_code == 200, response.text
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(response.json()["usage"]["cost"])
assert response.json()["usage"]["cost"] == pytest.approx(float(gateway_cost))