diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 1d88e133cf9..59444933989 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -48,7 +48,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, with_litellm_call_id, ) -from litellm.proxy.common_utils.token_count_log import rejected_token_count +from litellm.proxy.common_utils.token_count_log import answered_with_litellm_call_id, rejected_token_count from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.proxy.litellm_pre_call_utils import ( forwarded_anthropic_oauth_token, @@ -362,6 +362,7 @@ def _count_tokens_body_options(data: Mapping[str, JsonValue], auto_drop_params: ) async def count_tokens( request: Request, + fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # Used for auth ): """ @@ -384,9 +385,17 @@ async def count_tokens( Returns: {"input_tokens": } """ + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER)) + return await answered_with_litellm_call_id( + _count_tokens(request, user_api_key_dict, litellm_call_id), litellm_call_id, fastapi_response + ) + + +async def _count_tokens( + request: Request, user_api_key_dict: UserAPIKeyAuth, litellm_call_id: str +) -> JSONResponse | Mapping[str, object]: from litellm.proxy.proxy_server import count_request_tokens, llm_router, proxy_config - litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: request_data: Final = await read_request_body(request=request) data: Final[dict] = {**request_data} @@ -415,7 +424,9 @@ async def count_tokens( ) if isinstance(token_request, ProxyException): return rejected_token_count( - _anthropic_error_json_response(token_request, request), litellm_call_id, request_model + _anthropic_error_json_response(with_litellm_call_id(token_request, litellm_call_id), request), + litellm_call_id, + request_model, ) token_response: Final = await count_request_tokens( @@ -427,12 +438,16 @@ async def count_tokens( **_count_tokens_body_options(data, auto_drop_params), ) except HTTPException as e: - return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) + return _anthropic_error_json_response( + with_litellm_call_id(proxy_exception_from_http_exception(e, {}), litellm_call_id), request + ) except ProxyException as e: - return _anthropic_error_json_response(e, request) + return _anthropic_error_json_response(with_litellm_call_id(e, litellm_call_id), request) except Exception as e: log_llm_api_exception(e, litellm_call_id) - return _anthropic_error_json_response(proxy_exception_from_route_error(e), request) + return _anthropic_error_json_response( + with_litellm_call_id(proxy_exception_from_route_error(e), litellm_call_id), request + ) else: return {"input_tokens": token_response.total_tokens, **count_source_fields(token_response)} diff --git a/litellm/proxy/common_utils/token_count_log.py b/litellm/proxy/common_utils/token_count_log.py index 04c07a29117..e675f56bb98 100644 --- a/litellm/proxy/common_utils/token_count_log.py +++ b/litellm/proxy/common_utils/token_count_log.py @@ -6,10 +6,17 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final, TypeVar -from fastapi import status +from fastapi import HTTPException, Response, status from litellm._logging import verbose_proxy_logger -from litellm.proxy.common_utils.openai_error_payload import error_status_code +from litellm.proxy._types import ProxyException +from litellm.proxy.common_request_processing import log_llm_api_exception, proxy_exception_from_route_error +from litellm.proxy.common_utils.openai_error_payload import ( + LITELLM_CALL_ID_HEADER, + error_status_code, + headers_with_litellm_call_id, + with_litellm_call_id, +) from litellm.types.token_count import TokenCountSource from litellm.types.utils import TokenCountResponse @@ -28,6 +35,7 @@ def _logfmt_value(value: str | None) -> str: _Counted: Final = TypeVar("_Counted") _Rejection: Final = TypeVar("_Rejection") _ProviderCount: Final = TypeVar("_ProviderCount", bound=Mapping[str, object]) +_Answer: Final = TypeVar("_Answer") @dataclass(frozen=True, slots=True) @@ -131,3 +139,26 @@ async def logged_provider_token_count( ) return await _logged(count, _CountRequest(litellm_call_id, request_model, model_id), provider_fields) + + +async def answered_with_litellm_call_id( + answer: Awaitable[_Answer], litellm_call_id: str, fastapi_response: Response +) -> _Answer: + try: + answered: Final = await answer + except ProxyException as e: + raise with_litellm_call_id(e, litellm_call_id) + except HTTPException as e: + raise HTTPException( + status_code=e.status_code, + detail=e.detail, + headers=headers_with_litellm_call_id(e.headers, litellm_call_id), + ) from e + except Exception as e: + log_llm_api_exception(e, litellm_call_id) + raise with_litellm_call_id(proxy_exception_from_route_error(e), litellm_call_id) from e + if isinstance(answered, Response): + answered.headers.setdefault(LITELLM_CALL_ID_HEADER, litellm_call_id) + return answered + fastapi_response.headers[LITELLM_CALL_ID_HEADER] = litellm_call_id # rebind-ok: FastAPI's only way to add a header + return answered diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index c018f071b3c..0685ee1edba 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -22,7 +22,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg read_request_body, ) from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER -from litellm.proxy.common_utils.token_count_log import rejected_token_count +from litellm.proxy.common_utils.token_count_log import answered_with_litellm_call_id, rejected_token_count from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.types.llms.vertex_ai import TokenCountDetailsResponse @@ -327,7 +327,7 @@ async def google_stream_generate_content( dependencies=[Depends(user_api_key_auth)], response_model=TokenCountDetailsResponse, ) -async def google_count_tokens(request: Request, model_name: str): +async def google_count_tokens(request: Request, model_name: str, fastapi_response: Response): """ ```json return { @@ -342,10 +342,18 @@ async def google_count_tokens(request: Request, model_name: str): } ``` """ - from litellm.proxy.common_utils.http_parsing_utils import read_request_body - from litellm.proxy.proxy_server import token_counter as internal_token_counter - litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER)) + return await answered_with_litellm_call_id( + _google_count_tokens(request, model_name, litellm_call_id), litellm_call_id, fastapi_response + ) + + +async def _google_count_tokens( + request: Request, model_name: str, litellm_call_id: str +) -> JSONResponse | TokenCountDetailsResponse: + from litellm.proxy.common_utils.http_parsing_utils import read_request_body + from litellm.proxy.proxy_server import count_request_tokens + data: Final = await read_request_body(request=request) try: token_request: Final = _google_count_tokens_request(model_name, data) @@ -365,7 +373,7 @@ async def google_count_tokens(request: Request, model_name: str): ) try: - token_response: Final = await internal_token_counter( + token_response: Final = await count_request_tokens( request=token_request, call_endpoint=True, litellm_call_id=litellm_call_id ) except ProxyException as error: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 972ec18d487..ee23fdbfa7a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -107,6 +107,7 @@ from litellm.proxy.common_utils.sse_keepalive import ( wrap_passthrough_sse_bytes_with_keepalive_pings, ) from litellm.proxy.common_utils.token_count_log import ( + answered_with_litellm_call_id, deployment_model_id, logged_provider_token_count, rejected_token_count, @@ -1269,7 +1270,7 @@ async def handle_bedrock_count_tokens( fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, request_body: dict[str, object], -) -> dict[str, object]: +) -> Mapping[str, object]: """ Handle AWS Bedrock CountTokens API requests. @@ -1277,11 +1278,19 @@ async def handle_bedrock_count_tokens( - /v1/messages/count_tokens - /v1/messages/count-tokens """ + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER)) + return await answered_with_litellm_call_id( + _bedrock_count_tokens(user_api_key_dict, request_body, litellm_call_id), litellm_call_id, fastapi_response + ) + + +async def _bedrock_count_tokens( + user_api_key_dict: UserAPIKeyAuth, request_body: dict[str, object], litellm_call_id: str +) -> Mapping[str, object]: from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.proxy.proxy_server import llm_router - litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER)) try: # Initialize the handler handler: Final = BedrockCountTokensHandler() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 091e05610ba..28c7bb996e3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -559,6 +559,7 @@ from litellm.proxy.common_utils.timezone_utils import ( get_budget_reset_time, ) from litellm.proxy.common_utils.token_count_log import ( + answered_with_litellm_call_id, deployment_model_id, logged_token_count, rejected_token_count, @@ -15306,19 +15307,26 @@ def _system_message(system: object) -> ChatCompletionSystemMessage | None: ) async def token_counter( request: TokenCountRequest, + fastapi_response: Response, call_endpoint: bool = False, litellm_call_id: Annotated[str | None, Header(alias=LITELLM_CALL_ID_HEADER)] = None, ): """ Args: request: TokenCountRequest + fastapi_response: answers the x-litellm-call-id header the count's log line names call_endpoint: bool - When set to "True" it will call the token counting endpoint - e.g Anthropic or Google AI Studio Token Counting APIs. litellm_call_id: the x-litellm-call-id header, logged with the count so it joins the inference call's logs Returns: TokenCountResponse """ - return await count_request_tokens(request=request, call_endpoint=call_endpoint, litellm_call_id=litellm_call_id) + call_id: Final = resolve_litellm_call_id(litellm_call_id) + return await answered_with_litellm_call_id( + count_request_tokens(request=request, call_endpoint=call_endpoint, litellm_call_id=call_id), + call_id, + fastapi_response, + ) async def count_request_tokens( diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 9e642a02d49..223c6de876d 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -70,8 +70,9 @@ from litellm.proxy.common_utils.openai_error_payload import ( LITELLM_CALL_ID_HEADER, error_status_code, openai_error_type_for_status, + with_litellm_call_id, ) -from litellm.proxy.common_utils.token_count_log import rejected_token_count +from litellm.proxy.common_utils.token_count_log import answered_with_litellm_call_id, rejected_token_count from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decrypted_response_id @@ -1727,6 +1728,7 @@ def _responses_input_as_token_count_messages( ) async def responses_input_tokens( request: Request, + fastapi_response: Response, token_counter: _TokenCounter = _token_counter_dep, llm_router: Router | None = _llm_router_dep, user_api_key_dict: UserAPIKeyAuth = _user_api_key_auth_dep, @@ -1751,6 +1753,24 @@ async def responses_input_tokens( Returns: `{"object": "response.input_tokens", "input_tokens": }` """ litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER)) + return await answered_with_litellm_call_id( + _responses_input_tokens( + request, token_counter, llm_router, user_api_key_dict, response_id_security, managed_files, litellm_call_id + ), + litellm_call_id, + fastapi_response, + ) + + +async def _responses_input_tokens( + request: Request, + token_counter: _TokenCounter, + llm_router: Router | None, + user_api_key_dict: UserAPIKeyAuth, + response_id_security: ResponsesIDSecurity | None, + managed_files: _ManagedFileAccess | None, + litellm_call_id: str, +) -> JSONResponse | _ResponsesInputTokensResult: data: Final = await read_request_body(request=request) model_name: Final = data.get("model") input_value: Final = data.get("input") @@ -1833,7 +1853,7 @@ async def responses_input_tokens( except ProxyException as e: if e.type != _PROVIDER_COUNT_ERROR_TYPE: raise - return _provider_count_error_response(e, request) + return _provider_count_error_response(with_litellm_call_id(e, litellm_call_id), request) result: Final[_ResponsesInputTokensResult] = { "object": "response.input_tokens", "input_tokens": token_response.total_tokens, diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index 037b0a6c29a..bedd1430d3d 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -11,6 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from fastapi import Response from fastapi.testclient import TestClient from litellm import Router @@ -464,7 +465,9 @@ class TestErrorLogCarriesCallId: ), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), ): - response = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + response = await ep.count_tokens( + request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth() + ) assert response.status_code == 500 record = self._error_record(caplog) @@ -516,7 +519,9 @@ async def test_count_tokens_answers_400_for_malformed_tools_without_counting(mod patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts it is never reached patch.object(proxy_server, "llm_router", new=_claude_router()), # test-quality-ok: module global imported at call time; no injection seam on this route ): - response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + response: Final = await ep.count_tokens( + request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth() + ) assert response.status_code == 400 assert json.loads(response.body)["error"]["type"] == "invalid_request_error" @@ -539,7 +544,9 @@ async def test_count_tokens_counts_malformed_tool_entries_for_a_model_the_router patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts what it receives patch.object(proxy_server, "llm_router", new=_claude_router()), # test-quality-ok: module global imported at call time; no injection seam on this route ): - response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + response: Final = await ep.count_tokens( + request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth() + ) assert response == {"input_tokens": 7} assert counter.await_args.kwargs["request"].tools == _MALFORMED_TOOL_ENTRY @@ -572,7 +579,9 @@ async def test_count_tokens_drops_malformed_tools_when_drop_params_is_set( patch.object(ep, "read_request_body", new=AsyncMock(return_value=body)), # test-quality-ok: endpoint reads the body via a module function; no injection seam patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts what it receives ): - response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + response: Final = await ep.count_tokens( + request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth() + ) assert response == {"input_tokens": 7} counted: Final = counter.await_args.kwargs["request"] @@ -611,7 +620,9 @@ async def test_count_tokens_drops_malformed_tools_the_counted_deployment_drops_e patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts what it receives patch.object(proxy_server, "llm_router", new=router), # test-quality-ok: module global imported at call time; no injection seam on this route ): - response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + response: Final = await ep.count_tokens( + request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth() + ) assert response == {"input_tokens": 7} assert counter.await_args.kwargs["request"].tools is None @@ -781,7 +792,7 @@ def test_count_tokens_rejected_before_counting_logs_the_400_it_answers( ) -> None: response: Final = _count_on_local_model(monkeypatch, body) - assert response.status_code == 400, response.text + assert (response.status_code, response.headers.get("x-litellm-call-id")) == (400, "count-log-local"), response.text assert token_count_log_lines() == ( { "litellm_call_id": "count-log-local", @@ -795,6 +806,51 @@ def test_count_tokens_rejected_before_counting_logs_the_400_it_answers( ) +def _claude_count_reply(status_code: int) -> Callable[[httpx.Request], httpx.Response]: + reply: Final = ( + {"input_tokens": 77} + if status_code == 200 + else {"type": "error", "error": {"type": "invalid_request_error", "message": "invalid base64 data"}} + ) + return lambda request: httpx.Response(status_code, json=reply, request=request) + + +@pytest.mark.parametrize("client_call_id", [None, "count-echo-claude"]) +@pytest.mark.parametrize("provider_status", [200, 400]) +def test_count_tokens_answers_the_call_id_its_count_line_names( + monkeypatch: pytest.MonkeyPatch, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], + provider_status: int, + client_call_id: str | None, +) -> None: + response: Final = _count_on_claude( + monkeypatch, + _claude_count_reply(provider_status), + {} if client_call_id is None else {"x-litellm-call-id": client_call_id}, + ) + + (line,) = token_count_log_lines() + assert (response.status_code, response.headers.get("x-litellm-call-id")) == ( + provider_status, + line["litellm_call_id"], + ), response.text + assert client_call_id in (None, line["litellm_call_id"]) + + +def test_count_tokens_error_body_names_the_call_id_its_header_carries_when_opted_in( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", {"include_call_id_in_error_body": True}) + + response: Final = _count_on_claude(monkeypatch, _claude_count_reply(400), {}) + + (line,) = token_count_log_lines() + assert response.json()["error"]["litellm_call_id"] == response.headers.get("x-litellm-call-id") + assert response.headers.get("x-litellm-call-id") == line["litellm_call_id"] + + class TestEventLoggingBatchEndpoint: """Test the stubbed event logging batch endpoint""" diff --git a/tests/unit/proxy/common_utils/test_token_count_log.py b/tests/unit/proxy/common_utils/test_token_count_log.py index 94b976d88cf..5098d4994dd 100644 --- a/tests/unit/proxy/common_utils/test_token_count_log.py +++ b/tests/unit/proxy/common_utils/test_token_count_log.py @@ -1,9 +1,13 @@ import asyncio from collections.abc import Callable +from typing import Final import pytest +from fastapi import HTTPException, Response -from litellm.proxy.common_utils.token_count_log import logged_token_count +from litellm.proxy._types import ProxyException +from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER +from litellm.proxy.common_utils.token_count_log import answered_with_litellm_call_id, logged_token_count from litellm.types.utils import TokenCountResponse @@ -43,3 +47,39 @@ async def test_a_cancelled_count_writes_no_line( await logged_token_count(_count_ending_in(asyncio.CancelledError()), "count-log-cancelled", "count-log-gpt", None) assert token_count_log_lines() == () + + +@pytest.mark.asyncio +async def test_an_http_error_answers_its_status_and_detail_with_the_call_id_beside_its_own_headers() -> None: + rejection: Final = HTTPException(status_code=429, detail={"error": "slow down"}, headers={"retry-after": "3"}) + + with pytest.raises(HTTPException) as answered: + await answered_with_litellm_call_id(_count_ending_in(rejection), "count-echo-http", Response()) + + assert (answered.value.status_code, answered.value.detail, dict(answered.value.headers or {})) == ( + 429, + {"error": "slow down"}, + {"retry-after": "3", LITELLM_CALL_ID_HEADER: "count-echo-http"}, + ) + + +@pytest.mark.asyncio +async def test_an_unhandled_error_answers_a_500_with_the_call_id() -> None: + with pytest.raises(ProxyException) as answered: + await answered_with_litellm_call_id( + _count_ending_in(RuntimeError("tokenizer crashed")), "count-echo-crash", Response() + ) + + assert (answered.value.code, answered.value.headers) == ("500", {LITELLM_CALL_ID_HEADER: "count-echo-crash"}) + + +@pytest.mark.asyncio +async def test_a_cancelled_answer_stays_cancelled_and_names_no_call_id() -> None: + cancelled: Final = asyncio.CancelledError() + fastapi_response: Final = Response() + + with pytest.raises(asyncio.CancelledError) as raised: + await answered_with_litellm_call_id(_count_ending_in(cancelled), "count-echo-cancelled", fastapi_response) + + assert raised.value is cancelled + assert LITELLM_CALL_ID_HEADER not in fastapi_response.headers diff --git a/tests/unit/proxy/google_endpoints/test_endpoints.py b/tests/unit/proxy/google_endpoints/test_endpoints.py index 9dd2e267e64..8e1a39dd176 100644 --- a/tests/unit/proxy/google_endpoints/test_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_endpoints.py @@ -8,6 +8,7 @@ import pytest from dotenv import load_dotenv from pydantic import JsonValue, ValidationError from starlette.requests import Request +from starlette.responses import Response from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request, google_count_tokens from litellm.types.llms.vertex_ai import TokenCountDetailsResponse @@ -102,6 +103,7 @@ async def test_proxy_gemini_to_openai_like_model_token_counting(): } ), model_name="volcengine/foo", + fastapi_response=Response(), ) assert response.get("totalTokens") > 0 diff --git a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py index 4867ca9ab63..0509e6c68a3 100644 --- a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py @@ -172,7 +172,7 @@ def test_google_count_tokens_unchanged(): ) with patch( - "litellm.proxy.proxy_server.token_counter", + "litellm.proxy.proxy_server.count_request_tokens", new_callable=AsyncMock, return_value=fake_response, ): @@ -194,7 +194,7 @@ def test_google_count_tokens_unexpected_failure_uses_the_google_error_envelope() pytest.skip(f"Skipping test due to missing dependency: {e}") with patch( - "litellm.proxy.proxy_server.token_counter", + "litellm.proxy.proxy_server.count_request_tokens", new_callable=AsyncMock, side_effect=RuntimeError("tokenizer unavailable"), ): @@ -227,7 +227,7 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs() ) with patch( - "litellm.proxy.proxy_server.token_counter", + "litellm.proxy.proxy_server.count_request_tokens", new_callable=AsyncMock, return_value=fake_response, ): @@ -284,6 +284,7 @@ def test_google_count_tokens_logs_one_line_with_the_callers_call_id_model_and_co ) assert (response.status_code, response.json()["totalTokens"]) == (200, 77), response.text + assert response.headers.get("x-litellm-call-id") == "count-log-gemini" assert token_count_log_lines() == ( { "litellm_call_id": "count-log-gemini", @@ -318,3 +319,18 @@ def test_google_count_tokens_rejected_body_logs_the_400_it_answers( "status_code": 400, }, ) + + +@pytest.mark.parametrize("client_call_id", [None, "count-echo-gemini"]) +def test_google_count_tokens_rejected_body_answers_the_call_id_its_count_line_names( + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], client_call_id: str | None +) -> None: + response: Final = _build_test_client().post( + "/v1beta/models/count-log-claude:countTokens", + json={"contents": "not a list"}, + headers={} if client_call_id is None else {"x-litellm-call-id": client_call_id}, + ) + + (line,) = token_count_log_lines() + assert (response.status_code, response.headers.get("x-litellm-call-id")) == (400, line["litellm_call_id"]) + assert client_call_id in (None, line["litellm_call_id"]) diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 43fe5d4e786..0cd38eba1c6 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -7255,7 +7255,10 @@ _BEDROCK_COUNT_MODEL: Final = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1 async def _bedrock_native_count( - monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, aws_reply: httpx.Response + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + aws_reply: httpx.Response, + fastapi_response: Response, ) -> dict[str, object]: from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( handle_bedrock_count_tokens, @@ -7293,7 +7296,7 @@ async def _bedrock_native_count( return await handle_bedrock_count_tokens( endpoint=f"model/{_BEDROCK_COUNT_ALIAS}/count-tokens", request=request, - fastapi_response=Response(), + fastapi_response=fastapi_response, user_api_key_dict=MagicMock(), request_body={"model": _BEDROCK_COUNT_ALIAS, "messages": [{"role": "user", "content": "hi"}]}, ) @@ -7306,11 +7309,13 @@ async def test_bedrock_native_count_tokens_logs_the_count_against_the_router_mod respx_mock: respx.MockRouter, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], ) -> None: + fastapi_response: Final = Response() result: Final = await _bedrock_native_count( - monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12}) + monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12}), fastapi_response ) assert result == {"input_tokens": 12} + assert fastapi_response.headers.get("x-litellm-call-id") == "bedrock-native-count" assert len(respx_mock.calls) == 1 assert respx_mock.calls.last.request.url.path == "/model/anthropic.claude-haiku-4-5-20251001-v1:0/count-tokens" assert token_count_log_lines() == ( @@ -7335,9 +7340,12 @@ async def test_bedrock_native_count_tokens_logs_the_provider_status_when_the_cou ) -> None: rejected: Final = httpx.Response(400, json={"message": "The provided model doesn't support counting tokens."}) with pytest.raises(HTTPException) as exc_info: - await _bedrock_native_count(monkeypatch, respx_mock, rejected) + await _bedrock_native_count(monkeypatch, respx_mock, rejected, Response()) - assert exc_info.value.status_code == 400 + assert (exc_info.value.status_code, (exc_info.value.headers or {}).get("x-litellm-call-id")) == ( + 400, + "bedrock-native-count", + ) assert token_count_log_lines() == ( { "litellm_call_id": "bedrock-native-count", @@ -7391,6 +7399,33 @@ async def test_bedrock_native_count_tokens_without_a_model_logs_the_400_it_answe +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_without_a_model_answers_the_call_id_its_count_line_names( + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_bedrock_count_tokens, + ) + + request: Final = Request( + {"type": "http", "method": "POST", "path": f"/bedrock/model/{_BEDROCK_COUNT_ALIAS}/count-tokens", "headers": []} + ) + with pytest.raises(HTTPException) as exc_info: + await handle_bedrock_count_tokens( + endpoint=f"model/{_BEDROCK_COUNT_ALIAS}/count-tokens", + request=request, + fastapi_response=Response(), + user_api_key_dict=MagicMock(), + request_body={"messages": [{"role": "user", "content": "hi"}]}, + ) + + (line,) = token_count_log_lines() + assert (exc_info.value.status_code, (exc_info.value.headers or {}).get("x-litellm-call-id")) == ( + 400, + line["litellm_call_id"], + ) + + @pytest.mark.asyncio async def test_bedrock_native_count_tokens_with_a_region_only_model_answers_the_validation_error_and_logs_it( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index e3d514b2ea9..022969efde0 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -4830,6 +4830,7 @@ class TestResponsesInputTokens: ) assert (response.status_code, response.json().get("input_tokens")) == (200, 7), response.text + assert response.headers.get("x-litellm-call-id") == "count-log-responses" assert token_count_log_lines() == ( { "litellm_call_id": "count-log-responses", @@ -4863,6 +4864,48 @@ class TestResponsesInputTokens: }, ) + def test_count_rejected_before_counting_answers_the_call_id_its_count_line_names( + self, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] + ) -> None: + response, _ = self._post_input_tokens({"model": "count-log-gpt"}) + + (line,) = token_count_log_lines() + assert (response.status_code, response.headers.get("x-litellm-call-id")) == (400, line["litellm_call_id"]) + + @pytest.mark.parametrize( + ("raised", "status_code"), + [ + ( + ProxyException( + message="Token counting is disabled", type="token_counting_disabled", param="model", code=503 + ), + 503, + ), + ( + ProxyException( + message="LiteLLM does not support counting tokens for this model.", + type="token_counting_error", + param="model", + code=400, + ), + 400, + ), + (RuntimeError("tokenizer crashed"), 500), + ], + ids=["proxy-error", "provider-count-error", "crash"], + ) + def test_a_failed_count_answers_the_callers_call_id(self, raised: Exception, status_code: int) -> None: + response, _ = self._post_input_tokens( + {"model": "gpt-4o", "input": "hi"}, + counter=AsyncMock(side_effect=raised), + headers={"x-litellm-call-id": "count-echo-failed"}, + ) + + assert (response.status_code, response.headers.get("x-litellm-call-id")) == ( + status_code, + "count-echo-failed", + ), response.text + def test_a_managed_file_the_key_cannot_read_answers_403_without_a_count(self) -> None: from litellm.integrations.custom_logger import CustomLogger @@ -5149,6 +5192,21 @@ class TestResponsesInputTokens: } + def test_provider_count_error_body_names_the_call_id_its_header_carries_when_opted_in(self) -> None: + failing_counter: Final = AsyncMock( + side_effect=ProxyException( + message="LiteLLM does not support counting tokens for this model.", + type="token_counting_error", + param="model", + code="400", + ) + ) + with patch("litellm.proxy.proxy_server.general_settings", {"include_call_id_in_error_body": True}): + response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter) + + assert response.json()["error"]["litellm_call_id"] == response.headers.get("x-litellm-call-id") + + class TestGuardrailBlockedResponsesShape: """A pre_call block raises ModifyResponseException before any provider call. diff --git a/tests/unit/proxy/test_custom_tokenizer_bug.py b/tests/unit/proxy/test_custom_tokenizer_bug.py index 29625ea1a32..9b2c0498264 100644 --- a/tests/unit/proxy/test_custom_tokenizer_bug.py +++ b/tests/unit/proxy/test_custom_tokenizer_bug.py @@ -19,7 +19,7 @@ import litellm.proxy.proxy_server import litellm.utils from litellm import Router from litellm.proxy._types import ProxyException, TokenCountRequest -from litellm.proxy.proxy_server import token_counter +from litellm.proxy.proxy_server import count_request_tokens def _fake_hf_tokenizer(num_tokens: int) -> MagicMock: @@ -60,7 +60,7 @@ async def test_custom_tokenizer_from_model_info_is_used(monkeypatch): with patch.object(litellm.utils, "tokenizer_dispatch") as mock_tokenizer_cls: mock_tokenizer_cls.from_pretrained.return_value = _fake_hf_tokenizer(7) - response = await token_counter( + response = await count_request_tokens( request=TokenCountRequest( model="my-embedding-model", messages=[{"role": "user", "content": "Bonjour le monde"}], @@ -92,7 +92,7 @@ async def test_model_without_custom_tokenizer_uses_default(monkeypatch): monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", llm_router) with patch.object(litellm.utils, "tokenizer_dispatch") as mock_tokenizer_cls: - response = await token_counter( + response = await count_request_tokens( request=TokenCountRequest( model="gpt-4", messages=[{"role": "user", "content": "hello"}], @@ -133,7 +133,7 @@ async def test_custom_tokenizer_counts_a_vertex_partner_model_without_a_count_ro with patch.object(litellm.utils, "tokenizer_dispatch") as mock_tokenizer_cls: mock_tokenizer_cls.from_pretrained.return_value = _fake_hf_tokenizer(9) - response: Final = await token_counter( + response: Final = await count_request_tokens( request=TokenCountRequest(model="vertex-glm", messages=[{"role": "user", "content": "Return exactly OK."}]), call_endpoint=True, ) @@ -151,7 +151,7 @@ async def test_contents_only_count_of_a_vertex_partner_model_without_a_count_rou monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _vertex_glm_router_with_its_own_tokenizer()) with pytest.raises(ProxyException) as raised: - await token_counter( + await count_request_tokens( request=TokenCountRequest( model="vertex-glm", contents=[{"role": "user", "parts": [{"text": "Return exactly OK."}]}] ), diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index b21730c64f6..f3a7fd94b35 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -36,7 +36,7 @@ load_dotenv() # this file is to test litellm/proxy -from fastapi import Request +from fastapi import Request, Response from fastapi.responses import JSONResponse from starlette.types import Message @@ -78,7 +78,6 @@ from litellm.proxy.proxy_server import ( _provider_token_count, _try_provider_token_count, count_request_tokens, - token_counter, ) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted @@ -131,7 +130,7 @@ async def test_vLLM_token_counting(): setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - response = await token_counter( + response = await count_request_tokens( request=TokenCountRequest( model="special-alias", messages=[{"role": "user", "content": "hello"}], @@ -164,7 +163,7 @@ async def test_token_counting_model_not_in_model_list(): setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - response = await token_counter( + response = await count_request_tokens( request=TokenCountRequest( model="special-alias", messages=[{"role": "user", "content": "hello"}], @@ -197,7 +196,7 @@ async def test_gpt_token_counting(): setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - response = await token_counter( + response = await count_request_tokens( request=TokenCountRequest( model="gpt-4", messages=[{"role": "user", "content": "hello"}], @@ -240,7 +239,7 @@ async def test_token_counter_leaves_the_routers_deployment_without_the_credentia ) monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - response = await token_counter( + response = await count_request_tokens( request=TokenCountRequest(model="devstral", messages=[{"role": "user", "content": "hello"}]) ) @@ -271,10 +270,10 @@ async def test_token_counter_reads_an_oidc_api_key_only_when_it_calls_the_provid request = TokenCountRequest(model="claude", messages=[{"role": "user", "content": "hello"}]) if not call_endpoint: - assert (await token_counter(request=request, call_endpoint=call_endpoint)).total_tokens > 0 + assert (await count_request_tokens(request=request, call_endpoint=call_endpoint)).total_tokens > 0 return with pytest.raises(ProxyException) as excinfo: - await token_counter(request=request, call_endpoint=call_endpoint) + await count_request_tokens(request=request, call_endpoint=call_endpoint) assert excinfo.value.code == "401" assert str(tmp_path) not in excinfo.value.message @@ -389,7 +388,7 @@ async def test_anthropic_messages_count_tokens_endpoint(): try: # Call the endpoint - response = await count_tokens(mock_request, mock_user_api_key_dict) + response = await count_tokens(mock_request, Response(), mock_user_api_key_dict) # Verify response format matches Anthropic spec assert isinstance(response, dict) @@ -471,7 +470,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): try: # Call the endpoint - response = await count_tokens(mock_request, mock_user_api_key_dict) + response = await count_tokens(mock_request, Response(), mock_user_api_key_dict) # Verify response format matches Anthropic spec assert isinstance(response, dict) @@ -512,7 +511,7 @@ async def test_internal_token_counter_anthropic_provider_detection(): # Test with is_direct_request=False (simulating call from Anthropic endpoint) # The provider counter is engaged: its API error surfaces instead of a local count with pytest.raises(ProxyException) as exc_info: - await token_counter( + await count_request_tokens( request=TokenCountRequest( model="claude-test", messages=[{"role": "user", "content": "hello"}], @@ -536,7 +535,7 @@ async def test_internal_token_counter_anthropic_provider_detection(): setattr(litellm.proxy.proxy_server, "llm_router", llm_router) with pytest.raises(ProxyException) as missing_key: - await token_counter( + await count_request_tokens( request=TokenCountRequest( model="gpt-test", messages=[{"role": "user", "content": "hello"}], @@ -579,7 +578,7 @@ async def test_anthropic_endpoint_error_handling(): anthropic_endpoints.read_request_body = mock_read_request_body try: - response = await count_tokens(mock_request, mock_user_api_key_dict) + response = await count_tokens(mock_request, Response(), mock_user_api_key_dict) assert response.status_code == 400 assert json.loads(response.body) == { @@ -870,7 +869,7 @@ async def test_vertex_ai_anthropic_token_counting(): mock_handle_count_tokens.return_value = mock_token_response # Test with messages format and call_endpoint=True - response = await token_counter( + response = await count_request_tokens( request=TokenCountRequest( model="vertex_ai/claude-3-5-sonnet-20241022", messages=[ @@ -1085,7 +1084,7 @@ async def test_token_counter_httpx_status_error_raises_proxy_exception(): litellm.proxy.proxy_server.llm_router = mock_router with pytest.raises(ProxyException) as exc_info: - await token_counter( + await count_request_tokens( request=TokenCountRequest( model="claude-4-6-sonnet", messages=[{"role": "user", "content": "hello"}], @@ -1151,7 +1150,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter with pytest.raises(ProxyException) as exc_info: - await token_counter( + await count_request_tokens( request=TokenCountRequest( model="claude-bedrock", messages=[{"role": "user", "content": "hello"}], @@ -1216,7 +1215,7 @@ async def test_proxy_token_counter_error_raises_when_enabled(): with patch("litellm.token_counter") as mock_local_counter: with pytest.raises(ProxyException) as exc_info: - await token_counter( + await count_request_tokens( request=TokenCountRequest( model="claude-bedrock", messages=[{"role": "user", "content": "hello"}], @@ -1280,7 +1279,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): proxy_server.count_request_tokens = mock_token_counter_error try: - response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, Response(), mock_user_api_key_dict) assert isinstance(response, JSONResponse) assert response.status_code == 400 @@ -1339,7 +1338,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): proxy_server.count_request_tokens = mock_token_counter_error try: - response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, Response(), mock_user_api_key_dict) assert isinstance(response, JSONResponse) assert response.status_code == 403 @@ -1398,7 +1397,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): proxy_server.count_request_tokens = mock_token_counter_error try: - response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, Response(), mock_user_api_key_dict) assert isinstance(response, JSONResponse) assert response.status_code == 429 @@ -1474,7 +1473,7 @@ async def test_anthropic_count_tokens_drops_a_malformed_tool_result_text_when_dr route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens").mock( return_value=httpx.Response(200, json={"input_tokens": 9}) ) - response: Final = await anthropic_count_tokens(request, UserAPIKeyAuth()) + response: Final = await anthropic_count_tokens(request, Response(), UserAPIKeyAuth()) assert response["input_tokens"] == 9 sent: Final = json.loads(route.calls.last.request.content) @@ -1519,7 +1518,7 @@ async def _count_on_a_model_without_a_count_api( monkeypatch.setattr(proxy_server, "llm_router", router) monkeypatch.setattr(proxy_server, "proxy_config", proxy_server.ProxyConfig()) with respx.mock: - return await anthropic_count_tokens(request, UserAPIKeyAuth()) + return await anthropic_count_tokens(request, Response(), UserAPIKeyAuth()) @pytest.mark.asyncio @@ -1600,7 +1599,7 @@ async def test_anthropic_count_tokens_on_a_claude_deployment_drops_a_malformed_t route: Final = respx.post("https://api.anthropic.com/v1/messages/count_tokens").mock( return_value=httpx.Response(200, json={"input_tokens": 7}) ) - response: Final = await anthropic_count_tokens(request, UserAPIKeyAuth()) + response: Final = await anthropic_count_tokens(request, Response(), UserAPIKeyAuth()) assert response["input_tokens"] == 7 sent: Final = json.loads(route.calls.last.request.content) @@ -2090,8 +2089,8 @@ async def test_framed_estimate_that_cannot_read_the_body_answers_the_provider_er ) with pytest.raises(ProxyException) as raised: - await token_counter(request=TokenCountRequest(model="kimi", **unreadable_body), call_endpoint=True) - readable: Final = await token_counter( + await count_request_tokens(request=TokenCountRequest(model="kimi", **unreadable_body), call_endpoint=True) + readable: Final = await count_request_tokens( request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), call_endpoint=True, ) @@ -2138,8 +2137,8 @@ async def test_framed_estimate_counts_google_contents_like_the_same_conversation _framed_kimi(monkeypatch) google_request: Final = _google_count_tokens_request("kimi", {"contents": contents}) - from_google: Final = await token_counter(request=google_request, call_endpoint=True) - from_messages: Final = await token_counter( + from_google: Final = await count_request_tokens(request=google_request, call_endpoint=True) + from_messages: Final = await count_request_tokens( request=TokenCountRequest(model="kimi", messages=google_request.messages), call_endpoint=True ) @@ -2170,7 +2169,7 @@ async def test_framed_deployment_whose_custom_tokenizer_cannot_load_answers_the_ ) with pytest.raises(ProxyException) as raised: - await token_counter( + await count_request_tokens( request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), call_endpoint=True, ) @@ -2222,7 +2221,7 @@ def _framed_kimi_with_custom_tokenizer(monkeypatch: pytest.MonkeyPatch, identifi async def _count_status(model: str = "kimi") -> str: try: - response: Final = await token_counter( + response: Final = await count_request_tokens( request=TokenCountRequest(model=model, messages=[{"role": "user", "content": "What is new?"}]), call_endpoint=True, ) @@ -2486,8 +2485,8 @@ async def test_local_count_loads_the_tokenizer_again_after_a_failed_load(monkeyp request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]) with pytest.raises(OSError, match="hub unreachable"): - await token_counter(request=request) - retried: Final = await token_counter(request=request) + await count_request_tokens(request=request) + retried: Final = await count_request_tokens(request=request) assert retried.count_source == "local_estimate" and retried.total_tokens > 0 @@ -2517,15 +2516,15 @@ async def test_framed_count_loads_the_tokenizer_again_only_after_the_window_of_a monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_MAX_RUNNING", 1) request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]) - first: Final = await token_counter(request=request, call_endpoint=True) - second: Final = await token_counter(request=request, call_endpoint=True) + first: Final = await count_request_tokens(request=request, call_endpoint=True) + second: Final = await count_request_tokens(request=request, call_endpoint=True) assert (first.tokenizer_type, second.tokenizer_type) == tokenizer_types async def _fork_guard_traceback_frames() -> int: with pytest.raises(ForkedAfterNativeRuntimeStarted) as raised: - await token_counter( + await count_request_tokens( request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), call_endpoint=True, ) @@ -2564,7 +2563,7 @@ async def test_swallowed_estimate_error_is_logged_by_type_without_the_request_co verbose_proxy_logger.addHandler(caplog.handler) try: with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name), pytest.raises(ProxyException): - await token_counter( + await count_request_tokens( request=TokenCountRequest( model="kimi", messages=[ @@ -2609,9 +2608,9 @@ async def test_local_estimate_adds_the_configured_framing_once_per_request(monke ) messages: Final = [{"role": "system", "content": "be brief"}, {"role": "user", "content": "hello there"}] - plain: Final = await token_counter(request=TokenCountRequest(model="plain", messages=messages)) - framed: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) - framed_prompt: Final = await token_counter(request=TokenCountRequest(model="framed", prompt="hello there")) + plain: Final = await count_request_tokens(request=TokenCountRequest(model="plain", messages=messages)) + framed: Final = await count_request_tokens(request=TokenCountRequest(model="framed", messages=messages)) + framed_prompt: Final = await count_request_tokens(request=TokenCountRequest(model="framed", prompt="hello there")) assert plain.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) assert framed.total_tokens == plain.total_tokens + 100 @@ -2643,9 +2642,9 @@ async def test_framing_comes_from_the_row_the_deployment_provider_or_base_model_ messages: Final = [{"role": "user", "content": "hello there"}] framing: Final = cost_map_entry(kimi)["token_count_framing"] - prefixed: Final = await token_counter(request=TokenCountRequest(model="prefixed", messages=messages)) - split: Final = await token_counter(request=TokenCountRequest(model="split", messages=messages)) - aliased: Final = await token_counter(request=TokenCountRequest(model="aliased", messages=messages)) + prefixed: Final = await count_request_tokens(request=TokenCountRequest(model="prefixed", messages=messages)) + split: Final = await count_request_tokens(request=TokenCountRequest(model="split", messages=messages)) + aliased: Final = await count_request_tokens(request=TokenCountRequest(model="aliased", messages=messages)) assert prefixed.total_tokens == litellm.token_counter(model="system.ai.kimi-k3", messages=messages) + framing assert split.total_tokens == prefixed.total_tokens @@ -2750,7 +2749,9 @@ async def test_bedrock_kimi_k3_deployment_with_a_custom_tokenizer_keeps_counting monkeypatch.setattr(litellm.utils, "select_tokenizer", lambda model, custom_tokenizer=None: custom) messages: Final = [{"role": "user", "content": "What is new?"}] - counted: Final = await token_counter(request=TokenCountRequest(model="kimi", messages=messages), call_endpoint=True) + counted: Final = await count_request_tokens( + request=TokenCountRequest(model="kimi", messages=messages), call_endpoint=True + ) assert counted.total_tokens == ( litellm.token_counter(model=counted.model_used, messages=messages, custom_tokenizer=custom) @@ -2785,7 +2786,7 @@ async def test_bedrock_kimi_k3_without_its_tokenizer_adds_the_template_offsets_t ): _framed_kimi(monkeypatch) - counted: Final = await token_counter( + counted: Final = await count_request_tokens( request=TokenCountRequest(model="kimi", messages=messages, tools=tools), call_endpoint=True ) estimate: Final = litellm.token_counter( @@ -2814,7 +2815,7 @@ async def test_slow_cost_map_lookup_does_not_stall_other_requests_on_the_worker( monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None) monkeypatch.setattr(litellm, "get_model_info", slow_get_model_info) counting: Final = asyncio.create_task( - token_counter( + count_request_tokens( request=TokenCountRequest(model="slow-lookup-model", messages=[{"role": "user", "content": "hi"}]) ) ) @@ -2849,7 +2850,7 @@ async def test_uncountable_model_keeps_the_provider_error_when_local_counting_is monkeypatch.setattr(litellm, "disable_token_counter", True) with pytest.raises(ProxyException) as raised: - await token_counter( + await count_request_tokens( request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True ) @@ -2899,7 +2900,7 @@ async def test_configured_token_count_model_overrides_the_shipped_one_and_an_emp litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") ) - counted: Final = await token_counter( + counted: Final = await count_request_tokens( request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True ) @@ -2937,7 +2938,7 @@ async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geog litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") ) - await token_counter( + await count_request_tokens( request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True ) @@ -2995,7 +2996,7 @@ async def test_region_lookup_for_a_shipped_count_model_leaves_the_event_loop_fre monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) async def count() -> TokenCountResponse: - return await token_counter( + return await count_request_tokens( request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True, ) @@ -3023,7 +3024,7 @@ async def test_framing_still_applies_while_the_deployment_is_cooling_down(monkey ) monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) messages: Final = [{"role": "user", "content": "hello there"}] - warm: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + warm: Final = await count_request_tokens(request=TokenCountRequest(model="framed", messages=messages)) router.cooldown_cache.add_deployment_to_cooldown( model_id="framed-deployment", @@ -3031,7 +3032,7 @@ async def test_framing_still_applies_while_the_deployment_is_cooling_down(monkey exception_status=429, cooldown_time=60, ) - cooling: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + cooling: Final = await count_request_tokens(request=TokenCountRequest(model="framed", messages=messages)) assert warm.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + 100 assert cooling.total_tokens - 100 == litellm.token_counter(model="framed", messages=messages) @@ -3126,7 +3127,7 @@ async def _count_through_the_messages_endpoint( return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), ), ): - response: Final = await anthropic_count_tokens(mock_request, MagicMock()) + response: Final = await anthropic_count_tokens(mock_request, Response(), MagicMock()) return response, peer.request_bodies @@ -3447,10 +3448,10 @@ async def test_call_endpoint_counts_a_prompt_like_the_same_text_sent_as_a_user_m client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): - from_prompt: Final = await token_counter( + from_prompt: Final = await count_request_tokens( request=TokenCountRequest(model="gemini-counted", prompt=text), call_endpoint=True ) - from_message: Final = await token_counter( + from_message: Final = await count_request_tokens( request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": text}]), call_endpoint=True, ) @@ -3496,7 +3497,7 @@ async def test_call_endpoint_counts_a_prompt_on_a_vertex_ai_deployment(monkeypat with patch.object( VertexAIPartnerModelsTokenCounter, "handle_count_tokens_request", side_effect=_vertex_partner_word_count ): - from_prompt: Final = await token_counter( + from_prompt: Final = await count_request_tokens( request=TokenCountRequest(model="vertex-counted", prompt=text), call_endpoint=True ) @@ -3525,7 +3526,7 @@ async def test_prompt_next_to_a_conversation_answers_400_without_drop_params( patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client), pytest.raises(ProxyException) as rejected, ): - await token_counter( + await count_request_tokens( request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint ) @@ -3560,10 +3561,10 @@ async def test_drop_params_counts_a_conversation_without_the_prompt( client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): - with_prompt: Final = await token_counter( + with_prompt: Final = await count_request_tokens( request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint ) - conversation_only: Final = await token_counter( + conversation_only: Final = await count_request_tokens( request=TokenCountRequest(model="gemini-counted", **conversation), call_endpoint=call_endpoint ) @@ -3580,7 +3581,7 @@ async def test_call_endpoint_prompt_next_to_an_empty_conversation_counts_the_pro client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): - counted: Final = await token_counter( + counted: Final = await count_request_tokens( request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation), call_endpoint=True ) @@ -3593,8 +3594,10 @@ async def test_local_count_of_a_prompt_next_to_empty_messages_counts_the_prompt( ) -> None: monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) - with_empty: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX, messages=[])) - alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX)) + with_empty: Final = await count_request_tokens( + request=TokenCountRequest(model="gemini-counted", prompt=_FOX, messages=[]) + ) + alone: Final = await count_request_tokens(request=TokenCountRequest(model="gemini-counted", prompt=_FOX)) assert with_empty.error is False assert with_empty.total_tokens == alone.total_tokens > 0 @@ -3622,13 +3625,13 @@ async def test_local_count_of_a_prompt_with_system_or_tools_counts_them_like_the ) -> None: monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) - from_prompt: Final = await token_counter( + from_prompt: Final = await count_request_tokens( request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation, **chat_fields) ) - from_message: Final = await token_counter( + from_message: Final = await count_request_tokens( request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": _FOX}], **chat_fields) ) - prompt_alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX)) + prompt_alone: Final = await count_request_tokens(request=TokenCountRequest(model="gemini-counted", prompt=_FOX)) assert from_prompt.error is False assert from_prompt.total_tokens == from_message.total_tokens > prompt_alone.total_tokens @@ -3655,8 +3658,8 @@ async def test_call_endpoint_prompt_keeps_the_local_count_on_a_counter_that_need request: Final = TokenCountRequest(model="anthropic-counted", prompt="The quick brown fox jumps over the lazy dog.") with patch.object(anthropic_count_tokens_handler, "get_async_httpx_client", return_value=client): - with_endpoint: Final = await token_counter(request=request, call_endpoint=True) - local: Final = await token_counter(request=request, call_endpoint=False) + with_endpoint: Final = await count_request_tokens(request=request, call_endpoint=True) + local: Final = await count_request_tokens(request=request, call_endpoint=False) assert (with_endpoint.error, with_endpoint.total_tokens, with_endpoint.tokenizer_type) == ( False, @@ -4009,8 +4012,8 @@ async def test_token_counter_counts_a_custom_azure_deployment_name_with_its_base ) messages: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}] - named: Final = await token_counter(request=TokenCountRequest(model="named", messages=messages)) - served: Final = await token_counter(request=TokenCountRequest(model="served", messages=messages)) + named: Final = await count_request_tokens(request=TokenCountRequest(model="named", messages=messages)) + served: Final = await count_request_tokens(request=TokenCountRequest(model="served", messages=messages)) assert named.total_tokens == served.total_tokens @@ -4124,6 +4127,33 @@ def test_utils_token_counter_failed_count_logs_its_status_in_place_of_a_count( ) +@pytest.mark.parametrize( + ("body", "headers", "status_code"), + [ + (_COUNTED_HELLO, {}, 200), + (_COUNTED_HELLO, {"x-litellm-call-id": "count-echo-utils"}, 200), + ({"model": "count-log-gpt"}, {}, 400), + ({**_COUNTED_HELLO, "prompt": "and a prompt"}, {"x-litellm-call-id": "count-echo-utils"}, 400), + ], + ids=["counted", "counted-with-client-id", "no-messages", "failed-count-with-client-id"], +) +def test_utils_token_counter_answers_the_call_id_its_count_line_names( + monkeypatch: pytest.MonkeyPatch, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], + body: Mapping[str, JsonValue], + headers: Mapping[str, str], + status_code: int, +) -> None: + response: Final = _utils_token_counter(monkeypatch, body, headers) + + (line,) = token_count_log_lines() + assert (response.status_code, response.headers.get("x-litellm-call-id")) == ( + status_code, + line["litellm_call_id"], + ), response.text + assert headers.get("x-litellm-call-id") in (None, line["litellm_call_id"]) + + def test_utils_token_counter_json_line_keeps_its_own_fields_when_the_call_id_is_a_dict( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f114c1c73d2..d11bc4c4619 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -19068,6 +19068,7 @@ export interface paths { * Token Counter * @description Args: * request: TokenCountRequest + * fastapi_response: answers the x-litellm-call-id header the count's log line names * call_endpoint: bool - When set to "True" it will call the token counting endpoint - e.g Anthropic or Google AI Studio Token Counting APIs. * litellm_call_id: the x-litellm-call-id header, logged with the count so it joins the inference call's logs *