diff --git a/litellm/llms/vercel_ai_gateway/decisions/transformation.py b/litellm/llms/vercel_ai_gateway/decisions/transformation.py index 490b49bfff7..bd4051a04d1 100644 --- a/litellm/llms/vercel_ai_gateway/decisions/transformation.py +++ b/litellm/llms/vercel_ai_gateway/decisions/transformation.py @@ -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, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index f968e23f391..b2ca32fa2ea 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index f968e23f391..b2ca32fa2ea 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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, diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 4649e898454..d0d0067c16a 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -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) diff --git a/tests/unit/proxy/decisions_endpoints/test_endpoints.py b/tests/unit/proxy/decisions_endpoints/test_endpoints.py index e4e032360bf..79ebc78f098 100644 --- a/tests/unit/proxy/decisions_endpoints/test_endpoints.py +++ b/tests/unit/proxy/decisions_endpoints/test_endpoints.py @@ -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))