diff --git a/litellm/main.py b/litellm/main.py index 04b0963851a..ced63391fca 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -827,6 +827,25 @@ async def _sleep_for_timeout_async(timeout: float | str | httpx.Timeout): await asyncio.sleep(timeout.connect) +class _AdmissionReservation(BaseModel): + input_tokens: int | None = None + + +class _AdmissionMetadata(BaseModel): + user_api_key_budget_reservation: _AdmissionReservation | None = None + + +def admission_input_tokens(kwargs: Mapping[str, object]) -> int | None: + reservations: Final = ( + _AdmissionMetadata.model_validate(kwargs.get(key) or {}).user_api_key_budget_reservation + for key in ("litellm_metadata", "metadata") + ) + return next( + (reservation.input_tokens for reservation in reservations if reservation and reservation.input_tokens), + None, + ) + + def mock_completion( model: str, messages: list, @@ -838,6 +857,7 @@ def mock_completion( logging=None, custom_llm_provider=None, timeout: float | str | httpx.Timeout | None = None, + prompt_tokens: int | None = None, **kwargs, ): """ @@ -953,13 +973,16 @@ def mock_completion( ChatCompletionMessageToolCall(**tool_call) for tool_call in mock_tool_calls ] + usage_prompt_tokens: Final = ( + prompt_tokens if prompt_tokens is not None else DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT + ) setattr( model_response, "usage", Usage( - prompt_tokens=DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, + prompt_tokens=usage_prompt_tokens, completion_tokens=DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, - total_tokens=DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, + total_tokens=usage_prompt_tokens + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, ), ) @@ -5542,6 +5565,7 @@ def completion( custom_llm_provider=custom_llm_provider, mock_timeout=mock_timeout, timeout=timeout, + prompt_tokens=admission_input_tokens(cast(Mapping[str, object], kwargs)), ) ## RESPONSES API BRIDGE LOGIC ## - check if model has 'mode: responses' in litellm.model_cost map diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index b469cd70d76..61d82fb6158 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -318,6 +318,7 @@ async def reserve_budget_for_request( "entries": applied_entries, "finalized": False, "input_cost": min(float(input_cost or 0.0), reservation_cost), + "input_tokens": max(input_token_counts.values(), default=None), } diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index f272361062b..23eb297b47f 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -64,6 +64,16 @@ async def test_non_exempt_llm_route_still_reserves_budget(): assert reservation["reserved_cost"] > 0 +@pytest.mark.asyncio +async def test_reservation_carries_the_admission_input_token_count(): + reservation: Final = await _reserve("/v1/responses") + expected: Final = litellm.token_counter(model="gpt-4o", text="hello") + + assert reservation is not None + assert expected > 0 + assert reservation["input_tokens"] == expected + + ANTHROPIC_MESSAGES: Final = [{"role": "user", "content": "hello!!!"}] COUNT_TOKENS_REQUESTS: Final[tuple[tuple[str, dict[str, object]], ...]] = ( ("/v1/messages/count_tokens", {"model": "claude-sonnet-5", "messages": ANTHROPIC_MESSAGES}), diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index de6322f27ee..39ee5b57455 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2322,6 +2322,32 @@ def test_image_edit_merges_headers_and_extra_headers(): assert "extra_headers" not in handler_kwargs["image_edit_optional_request_params"] +@pytest.mark.parametrize("metadata_key", ("metadata", "litellm_metadata")) +def test_mock_completion_usage_reports_admission_input_tokens(metadata_key: str): + response = litellm.completion( + model="anthropic/claude-sonnet-5", + messages=[{"role": "user", "content": "hello"}], + mock_response="ok", + api_key="mock", + **{metadata_key: {"user_api_key_budget_reservation": {"reserved_cost": 1.0, "input_tokens": 51234}}}, + ) + + assert response.usage.prompt_tokens == 51234 + assert response.usage.total_tokens == 51234 + response.usage.completion_tokens + + +def test_mock_completion_usage_falls_back_to_default_without_admission_count(): + response = litellm.completion( + model="anthropic/claude-sonnet-5", + messages=[{"role": "user", "content": "hello"}], + mock_response="ok", + api_key="mock", + metadata={"user_api_key_budget_reservation": {"reserved_cost": 1.0}}, + ) + + assert response.usage.prompt_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT + + def test_mock_completion_stream_with_model_response(): """Test that mock_completion correctly handles stream=True with a ModelResponse as mock_response.""" from litellm import completion