From 36c7f10c8434894fea8d8ee715363c7b32cf1543 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Wed, 11 Nov 2026 14:26:47 -0800 Subject: [PATCH] fix(proxy): answer x-litellm-call-id on every token count response (#51037) * fix(proxy): answer x-litellm-call-id on every token count response The token count log line names a litellm_call_id, but count responses never carried it, so a client that sent no id of its own could not find its count's log line. Every count route (Anthropic Messages, Responses input_tokens, Gemini countTokens, Bedrock count-tokens, and /utils/token_counter) now answers the id on success and on errors, and a crash answers the JSON 500 the Anthropic and Gemini routes already used. * fix(proxy): keep the gateway's crash answers on count routes and stamp Responses opt-in bodies A crash on /utils/token_counter or /v1/responses/input_tokens goes back to the gateway's generic 500, or 503 when the database is down, now with the x-litellm-call-id header. Responses count rejections carry error.litellm_call_id when include_call_id_in_error_body is on * test(proxy): count through count_request_tokens in the token counter loop tests The token_counter route now also takes the response it sets x-litellm-call-id on, so these tests call the shared count function they exercise * fix(proxy): keep a count crash's own exception and log pre-count Responses failures The call id now reaches the crash handler through request.state instead of a wrapper exception, so logs and Sentry see the error the count raised. A Responses count that fails before counting writes its count line for any error, so its header names a line * fix(proxy): keep the token counter's injected args out of its API docs Tests the opt-in error body id on an Anthropic count the proxy rejects before counting * fix(proxy): keep a count route's own HTTP error when adding the call id The wrapper raised a new HTTPException with the same status, detail and headers, so a subclass lost its type and error tracking saw a new exception. It now adds the header to the original and raises it again --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../proxy/anthropic_endpoints/endpoints.py | 28 ++- litellm/proxy/common_utils/token_count_log.py | 31 ++- litellm/proxy/google_endpoints/endpoints.py | 20 +- .../llm_passthrough_endpoints.py | 16 +- litellm/proxy/proxy_server.py | 23 +- .../proxy/response_api_endpoints/endpoints.py | 81 +++++-- .../anthropic_endpoints/test_endpoints.py | 87 ++++++- .../common_utils/test_token_count_log.py | 52 ++++- .../proxy/google_endpoints/test_endpoints.py | 2 + .../test_google_api_endpoints.py | 22 +- .../test_llm_pass_through_endpoints.py | 45 +++- .../response_api_endpoints/test_endpoints.py | 134 ++++++++++- tests/unit/proxy/test_custom_tokenizer_bug.py | 10 +- ...test_proxy_server_endpoints_and_startup.py | 10 +- tests/unit/proxy/test_proxy_token_counter.py | 215 ++++++++++-------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 +- 16 files changed, 627 insertions(+), 157 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 636ba987ee8..ed3abbdc5b6 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -42,13 +42,14 @@ 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, error_status_code, litellm_call_id_headers, openai_error_param, 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.hooks.proxy_track_cost_callback import failure_row_response_cost from litellm.proxy.litellm_pre_call_utils import ( @@ -336,6 +337,7 @@ def _provider_context_management(response: TokenCountResponse) -> Mapping[str, o ) async def count_tokens( request: Request, + fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # Used for auth ): """ @@ -358,9 +360,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, request, 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} @@ -389,7 +399,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( @@ -401,13 +413,17 @@ 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: answered: Final = with_answered_status(e) log_llm_api_exception(answered, litellm_call_id) - return _anthropic_error_json_response(proxy_exception_from_route_error(answered), request) + return _anthropic_error_json_response( + with_litellm_call_id(proxy_exception_from_route_error(answered), litellm_call_id), request + ) else: return { "input_tokens": token_response.total_tokens, diff --git a/litellm/proxy/common_utils/token_count_log.py b/litellm/proxy/common_utils/token_count_log.py index 04c07a29117..83520ef1ade 100644 --- a/litellm/proxy/common_utils/token_count_log.py +++ b/litellm/proxy/common_utils/token_count_log.py @@ -6,10 +6,16 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final, TypeVar -from fastapi import status +from fastapi import HTTPException, Request, 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_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 +34,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 +138,23 @@ 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, request: Request, 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: + e.headers = headers_with_litellm_call_id(e.headers, litellm_call_id) + raise + except Exception: + request.state.litellm_call_id = litellm_call_id # rebind-ok: the crash handler reads the id from request.state + raise + 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..31b88e8110c 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, request, 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 f1a9e282bc9..321c921b5bd 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -133,6 +133,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, @@ -1526,7 +1527,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. @@ -1534,11 +1535,22 @@ 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, + request, + 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 b01dcfd2f03..fc4d4e65660 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -626,6 +626,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, @@ -2595,6 +2596,9 @@ async def otel_unhandled_exception_handler(request: Request, exc: Exception): raise exc response: Final = await _unhandled_exception_response(request, exc) response.headers.update(trace_response_headers(getattr(request.state, "parent_otel_span", None))) + litellm_call_id: Final = attribute_of(request.state, "litellm_call_id") + if isinstance(litellm_call_id, str): + response.headers.setdefault(LITELLM_CALL_ID_HEADER, litellm_call_id) return response @@ -16090,19 +16094,30 @@ def _system_message(system: object) -> ChatCompletionSystemMessage | None: ) async def token_counter( request: TokenCountRequest, + http_request: Request, + fastapi_response: Response, call_endpoint: bool = False, litellm_call_id: Annotated[str | None, Header(alias=LITELLM_CALL_ID_HEADER)] = None, ): """ + Returns: + TokenCountResponse + Args: request: TokenCountRequest 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 + \f + http_request: carries the call id to the error handler when the count fails + fastapi_response: answers the x-litellm-call-id header the count's log line names """ - 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, + http_request, + 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 369d167f623..b4eec96ecd9 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -11,6 +11,7 @@ from itertools import chain from types import MappingProxyType from typing import ( TYPE_CHECKING, + Annotated, Any, Final, Literal, @@ -80,8 +81,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 @@ -1575,6 +1577,12 @@ def _proxy_managed_files_hook() -> _ManagedFileAccess | None: _managed_files_hook_dep: Final = Depends(_proxy_managed_files_hook) +def _proxy_general_settings() -> Mapping[str, object]: + from litellm.proxy.proxy_server import general_settings_view + + return general_settings_view() + + async def _with_managed_file_id_mapping( data: Mapping[str, object], managed_files: _ManagedFileAccess | None, user_api_key_dict: UserAPIKeyAuth ) -> Mapping[str, object]: @@ -1607,7 +1615,9 @@ def _with_authorized_previous_response( return {**data, "previous_response_id": authorized_id, ADDRESSED_RESPONSE_ID_FIELD: previous_response_id} -def _responses_invalid_request_response(message: str, param: str | None, code: str | None) -> JSONResponse: +def _responses_invalid_request_response( + message: str, param: str | None, code: str | None, body_call_id: str | None +) -> JSONResponse: body: Final[_ResponsesApiErrorBody] = { "error": { "message": message, @@ -1616,7 +1626,7 @@ def _responses_invalid_request_response(message: str, param: str | None, code: s "code": code, } } - return JSONResponse(status_code=400, content=body) + return JSONResponse(status_code=400, content={"error": with_call_id({**body["error"]}, body_call_id)}) class _ProviderErrorMessage(TypedDict): @@ -1681,16 +1691,17 @@ def _provider_count_error_body(raw_message: str, status_code: int) -> _Responses } -def _provider_count_error_response(exc: ProxyException, request: Request) -> JSONResponse: +def _provider_count_error_response( + exc: ProxyException, request: Request, general_settings: Mapping[str, object] +) -> JSONResponse: from litellm.proxy.proxy_server import ( _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; a route that answers an error itself must stamp the OTel server span like the global ProxyException handler does - general_settings_view, ) status_code: Final = error_status_code(exc, 500) _close_dangling_otel_server_span(request, status_code, exc=exc) body: Final = _provider_count_error_body(exc.message, status_code) - call_id: Final = error_body_call_id(general_settings_view(), exc.headers.get(LITELLM_CALL_ID_HEADER)) + call_id: Final = error_body_call_id(general_settings, exc.headers.get(LITELLM_CALL_ID_HEADER)) return JSONResponse( status_code=status_code, content={"error": with_call_id({**body["error"]}, call_id)}, @@ -1698,11 +1709,12 @@ def _provider_count_error_response(exc: ProxyException, request: Request) -> JSO ) -def _missing_responses_param_response(param: str) -> JSONResponse: +def _missing_responses_param_response(param: str, body_call_id: str | None) -> JSONResponse: return _responses_invalid_request_response( message=f"Missing required parameter: '{param}'.", param=param, code="missing_required_parameter", + body_call_id=body_call_id, ) @@ -1744,6 +1756,8 @@ def _responses_input_as_token_count_messages( ) async def responses_input_tokens( request: Request, + fastapi_response: Response, + general_settings: Annotated[Mapping[str, object], Depends(_proxy_general_settings)], token_counter: _TokenCounter = _token_counter_dep, llm_router: Router | None = _llm_router_dep, user_api_key_dict: UserAPIKeyAuth = _user_api_key_auth_dep, @@ -1768,19 +1782,50 @@ 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, + general_settings, + ), + litellm_call_id, + request, + 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, + general_settings: Mapping[str, object], +) -> JSONResponse | _ResponsesInputTokensResult: + body_call_id: Final = error_body_call_id(general_settings, litellm_call_id) data: Final = await read_request_body(request=request) model_name: Final = data.get("model") input_value: Final = data.get("input") if not isinstance(model_name, str) or not model_name: - return rejected_token_count(_missing_responses_param_response("model"), litellm_call_id, None) + return rejected_token_count(_missing_responses_param_response("model", body_call_id), litellm_call_id, None) if input_value is None: - return rejected_token_count(_missing_responses_param_response("input"), litellm_call_id, model_name) + return rejected_token_count( + _missing_responses_param_response("input", body_call_id), litellm_call_id, model_name + ) if isinstance(input_value, (str, list)) and not input_value: return rejected_token_count( _responses_invalid_request_response( message="""One of "input" or "previous_response_id" or 'prompt' or 'conversation' must be provided.""", param=None, code="missing_required_parameter", + body_call_id=body_call_id, ), litellm_call_id, model_name, @@ -1792,7 +1837,9 @@ async def responses_input_tokens( token_fields: Final = resolve_counted_tool_fields(counted_fields, llm_router, user_api_key_dict.team_id) if isinstance(token_fields, ProxyException): return rejected_token_count( - _responses_invalid_request_response(message=token_fields.message, param=token_fields.param, code=None), + _responses_invalid_request_response( + message=token_fields.message, param=token_fields.param, code=None, body_call_id=body_call_id + ), litellm_call_id, model_name, ) @@ -1807,7 +1854,10 @@ async def responses_input_tokens( except ValidationError as e: return rejected_token_count( _responses_invalid_request_response( - message=first_validation_error(e), param=str(e.errors(include_url=False)[0]["loc"][0]), code=None + message=first_validation_error(e), + param=str(e.errors(include_url=False)[0]["loc"][0]), + code=None, + body_call_id=body_call_id, ), litellm_call_id, model_name, @@ -1815,7 +1865,7 @@ async def responses_input_tokens( except Exception as e: return rejected_token_count( _responses_invalid_request_response( - message=f"Invalid request for token counting: {e}", param=None, code=None + message=f"Invalid request for token counting: {e}", param=None, code=None, body_call_id=body_call_id ), litellm_call_id, model_name, @@ -1835,8 +1885,9 @@ async def responses_input_tokens( managed_files, user_api_key_dict, ) - except ProxyException as e: - raise rejected_token_count(e, litellm_call_id, model_name) + except Exception as e: + rejected_token_count(e, litellm_call_id, model_name) + raise try: token_response: Final = await token_counter( request=token_request, @@ -1847,7 +1898,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, general_settings) 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 072de5c0eed..f8706efb4bd 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -557,7 +557,9 @@ class TestProxyExceptionAnthropicEnvelope: proxy_server, "count_request_tokens", new=AsyncMock(side_effect=unavailable) ), ): - 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() + ) body: Final = json.loads(response.body) assert response.status_code == expected_status @@ -840,7 +842,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) @@ -883,7 +887,9 @@ class TestCountTokensCompactedHistory: 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, "llm_router", None), # test-quality-ok: module global; with no router the count runs on the local tokenizer, the path that raised ): - return await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + return await ep.count_tokens( + request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth() + ) counted = await count(compacted) @@ -934,7 +940,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" @@ -957,7 +965,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 @@ -990,7 +1000,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"] @@ -1017,7 +1029,9 @@ async def test_count_tokens_keeps_claude_code_malformed_tools_when_operator_sets patch.object(proxy_server, "llm_router", new=_claude_router()), # test-quality-ok: module global imported at call time; no injection seam on this route patch.object(proxy_server, "proxy_config", new=operator_config), # 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" @@ -1055,7 +1069,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 @@ -1246,7 +1262,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", @@ -1303,6 +1319,59 @@ def test_count_tokens_on_bedrock_answers_400_for_a_message_without_content_witho assert len(respx_mock.calls) == 0 +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"]) + + +@pytest.mark.parametrize( + ("provider_status", "extra_body"), + [(400, MappingProxyType({})), (200, MappingProxyType({"tools": 5}))], + ids=["provider_rejection", "malformed_tools_rejection"], +) +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], ...]], + provider_status: int, + extra_body: Mapping[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(provider_status), {}, extra_body) + + (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..b34bb55eb82 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,12 @@ import asyncio from collections.abc import Callable +from typing import Final import pytest +from fastapi import HTTPException, Request, Response -from litellm.proxy.common_utils.token_count_log import logged_token_count +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 @@ -15,6 +18,10 @@ async def _count_ending_in(error: BaseException) -> TokenCountResponse: raise error +def _count_request() -> Request: + return Request({"type": "http", "method": "POST", "path": "/utils/token_counter", "headers": []}) + + @pytest.mark.asyncio async def test_a_count_ending_in_a_base_exception_logs_a_500( token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], @@ -43,3 +50,46 @@ 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", _count_request(), 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"}, + ) + + +class _CountTooLarge(HTTPException): + pass + + +@pytest.mark.asyncio +async def test_an_http_error_reaches_its_handler_as_the_same_error_with_the_call_id() -> None: + too_large: Final = _CountTooLarge(status_code=413, detail="count body too large") + + with pytest.raises(_CountTooLarge) as answered: + await answered_with_litellm_call_id(_count_ending_in(too_large), "count-echo-same", _count_request(), Response()) + + assert answered.value is too_large + assert dict(answered.value.headers or {}) == {LITELLM_CALL_ID_HEADER: "count-echo-same"} + + +@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", _count_request(), 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 8e88fd7449a..7f118fb3186 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 @@ -8250,6 +8250,7 @@ async def _bedrock_native_count( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, aws_reply: httpx.Response, + fastapi_response: Response, request_body: Mapping[str, object] = _BEDROCK_COUNT_BODY, ) -> dict[str, object]: from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( @@ -8288,7 +8289,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=dict(request_body), ) @@ -8301,9 +8302,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: - result: Final = await _bedrock_native_count(monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12})) + fastapi_response: Final = Response() + result: Final = await _bedrock_native_count( + 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() == ( @@ -8328,9 +8333,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", @@ -8363,7 +8371,7 @@ async def test_bedrock_native_count_tokens_answers_400_for_a_malformed_body_with ) -> None: with pytest.raises(HTTPException) as exc_info: await _bedrock_native_count( - monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12}), request_body + monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12}), Response(), request_body ) assert (exc_info.value.status_code, exc_info.value.detail) == (400, {"error": problem}) @@ -8458,6 +8466,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 e4e4b27e1c2..0a28bb507d2 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -5416,10 +5416,13 @@ class TestResponsesInputTokens: response_id_security: object | None = None, managed_files: object | None = None, headers: Mapping[str, str] = EMPTY_MAPPING, + general_settings: Mapping[str, object] | None = None, + raise_server_exceptions: bool = False, ) -> tuple[Response, AsyncMock]: from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.response_api_endpoints.endpoints import ( + _proxy_general_settings, _proxy_llm_router, _proxy_managed_files_hook, _proxy_response_id_security, @@ -5446,8 +5449,10 @@ class TestResponsesInputTokens: app.dependency_overrides[_proxy_response_id_security] = lambda: response_id_security if managed_files is not None: app.dependency_overrides[_proxy_managed_files_hook] = lambda: managed_files + if general_settings is not None: + app.dependency_overrides[_proxy_general_settings] = lambda: general_settings try: - client = TestClient(app) + client = TestClient(app, raise_server_exceptions=raise_server_exceptions) response = client.post(path, json=body, headers={"Authorization": "Bearer sk-9876", **headers}) return response, token_counter_mock finally: @@ -5456,6 +5461,7 @@ class TestResponsesInputTokens: app.dependency_overrides.pop(_proxy_llm_router, None) app.dependency_overrides.pop(_proxy_response_id_security, None) app.dependency_overrides.pop(_proxy_managed_files_hook, None) + app.dependency_overrides.pop(_proxy_general_settings, None) def test_string_input_returns_openai_input_tokens_shape(self): response, counter = self._post_input_tokens({"model": "gpt-4o", "input": "Hello, how are you?"}) @@ -5738,6 +5744,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", @@ -5771,6 +5778,81 @@ 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, + ), + ], + ids=["proxy-error", "provider-count-error"], + ) + 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_crashed_count_answers_the_gateways_internal_error_with_the_callers_call_id(self) -> None: + response, _ = self._post_input_tokens( + {"model": "gpt-4o", "input": "hi"}, + counter=AsyncMock(side_effect=RuntimeError("tokenizer crashed reading /etc/litellm/config.yaml")), + headers={"x-litellm-call-id": "count-echo-crashed"}, + ) + + assert (response.status_code, response.headers.get("x-litellm-call-id")) == (500, "count-echo-crashed") + assert response.json() == {"error": {"message": "Internal server error", "type": "internal_server_error"}} + + def test_a_crashed_count_reaches_error_tracking_as_the_error_it_raised(self) -> None: + crash: Final = RuntimeError("tokenizer crashed") + + with pytest.raises(RuntimeError) as raised: + self._post_input_tokens( + {"model": "gpt-4o", "input": "hi"}, + counter=AsyncMock(side_effect=crash), + raise_server_exceptions=True, + ) + + assert raised.value is crash + + def test_a_count_during_a_database_outage_answers_503_with_the_callers_call_id(self) -> None: + from prisma.errors import ClientNotConnectedError + + response, _ = self._post_input_tokens( + {"model": "gpt-4o", "input": "hi"}, + counter=AsyncMock(side_effect=ClientNotConnectedError()), + headers={"x-litellm-call-id": "count-echo-db-down"}, + ) + + assert (response.status_code, response.headers.get("x-litellm-call-id")) == (503, "count-echo-db-down") + assert response.json()["error"]["type"] == "no_db_connection" + def test_a_managed_file_the_key_cannot_read_answers_403_without_a_count(self) -> None: from litellm.integrations.custom_logger import CustomLogger @@ -5792,6 +5874,28 @@ class TestResponsesInputTokens: assert response.json()["error"]["message"] == denied counter.assert_not_awaited() + def test_a_managed_file_lookup_crash_answers_the_call_id_its_count_line_names( + self, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] + ) -> None: + from litellm.integrations.custom_logger import CustomLogger + + class ManagedFiles(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + raise RuntimeError("managed file table unreadable") + + response, counter = self._post_input_tokens( + { + "model": "count-log-gpt", + "input": [{"role": "user", "content": [{"type": "input_file", "file_id": "litellm-managed-file"}]}], + }, + managed_files=ManagedFiles(), + ) + + (line,) = token_count_log_lines() + assert (response.status_code, response.headers.get("x-litellm-call-id")) == (500, line["litellm_call_id"]) + assert line["status_code"] == 500 + counter.assert_not_awaited() + def test_malformed_tools_follow_the_counted_deployment_even_when_fallbacks_are_configured(self) -> None: router: Final = litellm.Router( model_list=[ @@ -6057,6 +6161,34 @@ 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") + + @pytest.mark.parametrize( + "body", + [{"input": "hi"}, {"model": "gpt-4o"}, {"model": "gpt-4o", "input": ""}], + ) + def test_count_rejection_body_names_the_call_id_its_header_carries_when_opted_in( + self, body: Mapping[str, object] + ) -> None: + response, counter = self._post_input_tokens(dict(body), general_settings={"include_call_id_in_error_body": True}) + + counter.assert_not_awaited() + assert response.status_code == 400, response.text + 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_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 5b423421458..4f96af7f02c 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -15778,7 +15778,7 @@ async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_coun warm_tokenizer("claude-fable-5") response, took, lags = await timed_with_loop_lags( - lambda: proxy_server_module.token_counter(TokenCountRequest(model="claude-fable-5", prompt=text * 100)) + lambda: proxy_server_module.count_request_tokens(TokenCountRequest(model="claude-fable-5", prompt=text * 100)) ) assert response.total_tokens > 0 @@ -15819,7 +15819,7 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp litellm.utils._select_custom_tokenizer_helper.cache_clear() try: - response: Final = await proxy_server_module.token_counter( + response: Final = await proxy_server_module.count_request_tokens( TokenCountRequest(model="self-hosted", prompt="count me off the loop") ) assert loader_threads.get_nowait() is not loop_thread @@ -15860,7 +15860,7 @@ async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revisi litellm.utils._select_custom_tokenizer_helper.cache_clear() try: responses: Final = [ - await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me once")) + await proxy_server_module.count_request_tokens(TokenCountRequest(model="self-hosted", prompt="count me once")) for _ in range(3) ] assert from_pretrained.call_args_list == [mock.call("my-org/tokenizer", revision="main", token=None)] @@ -15868,8 +15868,8 @@ async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revisi assert len({response.total_tokens for response in responses}) == 1 assert responses[0].total_tokens > 0 - await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-pinned", prompt="count me once")) - await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-private", prompt="count me once")) + await proxy_server_module.count_request_tokens(TokenCountRequest(model="self-hosted-pinned", prompt="count me once")) + await proxy_server_module.count_request_tokens(TokenCountRequest(model="self-hosted-private", prompt="count me once")) assert from_pretrained.call_args_list == [ mock.call("my-org/tokenizer", revision="main", token=None), mock.call("my-org/tokenizer", revision="v2", token=None), diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 34868606012..bbd3ce8c78d 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -38,7 +38,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 @@ -83,7 +83,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 @@ -139,7 +138,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"}], @@ -172,7 +171,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"}], @@ -209,8 +208,8 @@ async def test_asking_the_provider_endpoint_leaves_an_openrouter_deployments_cou model="or-gpt", messages=[{"role": "user", "content": "Count these tokens, ünïcödé 你好 🙂"}] ) - via_provider_endpoint: Final = await token_counter(request=request, call_endpoint=True) - counted_locally: Final = await token_counter(request=request, call_endpoint=False) + via_provider_endpoint: Final = await count_request_tokens(request=request, call_endpoint=True) + counted_locally: Final = await count_request_tokens(request=request, call_endpoint=False) assert (via_provider_endpoint.total_tokens, via_provider_endpoint.model_used) == ( counted_locally.total_tokens, @@ -238,7 +237,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"}], @@ -281,7 +280,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"}]) ) @@ -312,10 +311,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 @@ -343,7 +342,7 @@ async def test_token_counter_answers_503_when_the_metadata_server_behind_an_oidc 500, text="metadata server unavailable" ) with pytest.raises(ProxyException) as excinfo: - await token_counter( + await count_request_tokens( request=TokenCountRequest(model="claude", messages=[{"role": "user", "content": "hello"}]), call_endpoint=True, ) @@ -463,7 +462,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) @@ -545,7 +544,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) @@ -586,7 +585,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"}], @@ -610,7 +609,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"}], @@ -653,7 +652,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) == { @@ -947,7 +946,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=[ @@ -1161,7 +1160,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"}], @@ -1227,7 +1226,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"}], @@ -1292,7 +1291,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"}], @@ -1356,7 +1355,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 @@ -1415,7 +1414,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 @@ -1474,7 +1473,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 @@ -1550,7 +1549,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) @@ -1595,7 +1594,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 @@ -1676,7 +1675,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) @@ -2191,8 +2190,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, ) @@ -2223,7 +2222,7 @@ async def test_framed_estimate_counts_a_json_schema_shorthand_like_its_explicit_ ) async def count(schema: object) -> TokenCountResponse: - return await token_counter( + return await count_request_tokens( request=TokenCountRequest( model="kimi", messages=[{"role": "user", "content": "What is new?"}], @@ -2277,8 +2276,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 ) @@ -2309,7 +2308,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, ) @@ -2361,7 +2360,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, ) @@ -2625,8 +2624,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 @@ -2656,15 +2655,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, ) @@ -2703,7 +2702,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=[ @@ -2748,9 +2747,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 @@ -2782,9 +2781,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 @@ -2964,7 +2963,7 @@ async def test_local_estimate_matches_the_billed_prompt_with_the_shipped_framing ), ) - counted: Final = await token_counter( + counted: Final = await count_request_tokens( request=TokenCountRequest( model="billed", messages=[{"role": "user", "content": "What is the weather in Paris?"}], @@ -3363,7 +3362,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) @@ -3398,7 +3399,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( @@ -3427,7 +3428,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"}]) ) ) @@ -3462,7 +3463,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 ) @@ -3512,7 +3513,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 ) @@ -3550,7 +3551,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 ) @@ -3608,7 +3609,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, ) @@ -3636,7 +3637,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", @@ -3644,7 +3645,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) @@ -3738,7 +3739,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 @@ -4098,7 +4099,7 @@ async def _messages_endpoint_count(router: Router, request_body: Mapping[str, Js patch.object(proxy_server, "llm_router", router), patch.object(anthropic_endpoints, "read_request_body", AsyncMock(return_value=dict(request_body))), ): - return await anthropic_count_tokens(mock_request, MagicMock()) + return await anthropic_count_tokens(mock_request, Response(), MagicMock()) class _MessagesCountReply(BaseModel): @@ -4651,10 +4652,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, ) @@ -4700,7 +4701,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 ) @@ -4729,7 +4730,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 ) @@ -4764,10 +4765,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 ) @@ -4784,7 +4785,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 ) @@ -4797,8 +4798,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 @@ -4826,13 +4829,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 @@ -4903,8 +4906,8 @@ async def test_call_endpoint_sends_a_prompt_to_the_provider_as_one_user_message( ) with _count_apis_answering(count_api): - from_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) - from_message: Final = await token_counter(request=_FOX_AS_MESSAGE, call_endpoint=True) + from_prompt: Final = await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=True) + from_message: Final = await count_request_tokens(request=_FOX_AS_MESSAGE, call_endpoint=True) assert (from_prompt.total_tokens, from_prompt.error) == (41, False) assert from_prompt.count_source != "local_estimate" @@ -4925,9 +4928,9 @@ async def test_call_endpoint_answers_a_provider_error_on_a_prompt_like_on_the_sa with _count_apis_answering(count_api): with pytest.raises(ProxyException) as from_prompt: - await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=True) with pytest.raises(ProxyException) as from_message: - await token_counter(request=_FOX_AS_MESSAGE, call_endpoint=True) + await count_request_tokens(request=_FOX_AS_MESSAGE, call_endpoint=True) assert from_prompt.value.code == str(status) assert f"scripted {status}" in from_prompt.value.message @@ -4944,11 +4947,11 @@ async def test_framed_stand_in_counts_a_prompt_like_the_same_text_sent_as_a_user ) -> None: _framed_kimi(monkeypatch) - from_prompt: Final = await token_counter(request=TokenCountRequest(model="kimi", prompt=_FOX), call_endpoint=True) - from_message: Final = await token_counter( + from_prompt: Final = await count_request_tokens(request=TokenCountRequest(model="kimi", prompt=_FOX), call_endpoint=True) + from_message: Final = await count_request_tokens( request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": _FOX}]), call_endpoint=True ) - local_prompt: Final = await token_counter(request=TokenCountRequest(model="kimi", prompt=_FOX)) + local_prompt: Final = await count_request_tokens(request=TokenCountRequest(model="kimi", prompt=_FOX)) assert (from_prompt.count_source, from_prompt.error) == ("local_estimate", False) assert from_prompt.total_tokens == from_message.total_tokens > local_prompt.total_tokens @@ -4970,9 +4973,9 @@ async def test_call_endpoint_counts_a_prompt_on_a_provider_without_a_count_api_a return httpx.Response(200, json={"input_tokens": 41}, request=request) with _count_apis_answering(count_api): - from_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) - from_message: Final = await token_counter(request=_FOX_AS_MESSAGE, call_endpoint=True) - local_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=False) + from_prompt: Final = await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=True) + from_message: Final = await count_request_tokens(request=_FOX_AS_MESSAGE, call_endpoint=True) + local_prompt: Final = await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=False) assert sent == [] assert (from_prompt.count_source, from_prompt.error) == ("local_estimate", False) @@ -5008,8 +5011,8 @@ async def test_call_endpoint_counts_a_prompt_on_a_non_chat_deployment_as_raw_tex return httpx.Response(200, json={"input_tokens": 41}, request=request) with _count_apis_answering(count_api): - from_endpoint: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) - local: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=False) + from_endpoint: Final = await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=True) + local: Final = await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=False) assert sent == [] assert (from_endpoint.total_tokens, from_endpoint.error) == (local.total_tokens, False) @@ -5044,7 +5047,7 @@ async def test_call_endpoint_sends_a_prompt_on_a_chat_mode_deployment_to_the_pro return httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 41}, request=request) with _count_apis_answering(count_api): - from_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + from_prompt: Final = await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=True) assert (from_prompt.total_tokens, from_prompt.error) == (41, False) assert len(sent) == 1 @@ -5061,8 +5064,8 @@ async def test_call_endpoint_counts_a_prompt_for_a_model_without_a_deployment_as ) request: Final = TokenCountRequest(model="deepseek/deepseek-v4-flash", prompt=_FOX) - from_endpoint: Final = await token_counter(request=request, call_endpoint=True) - local: Final = await token_counter(request=request, call_endpoint=False) + from_endpoint: Final = await count_request_tokens(request=request, call_endpoint=True) + local: Final = await count_request_tokens(request=request, call_endpoint=False) assert (from_endpoint.total_tokens, from_endpoint.error) == (local.total_tokens, False) @@ -5102,7 +5105,7 @@ async def test_call_endpoint_answers_a_timeout_when_a_stalled_count_api_outlasts with pytest.raises(ProxyException) as error: async with asyncio.timeout(10): - await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=True) assert error.value.code == "408" @@ -5144,7 +5147,7 @@ async def test_call_endpoint_names_the_status_of_a_count_api_error_with_an_empty with _count_apis_answering(lambda request: httpx.Response(404, request=request)): with pytest.raises(ProxyException) as error: - await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=True) assert error.value.code == "404" assert "404" in error.value.message @@ -5494,8 +5497,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 @@ -5527,7 +5530,9 @@ def _utils_token_counter( ), ) monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-test")) - return TestClient(app).post("/utils/token_counter", json=dict(body), headers=dict(headers)) + return TestClient(app, raise_server_exceptions=False).post( + "/utils/token_counter", json=dict(body), headers=dict(headers) + ) def test_utils_token_counter_logs_one_line_with_the_callers_call_id_model_and_count( @@ -5609,6 +5614,38 @@ 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), + ( + {"model": "count-log-gpt", "messages": [{"role": "user", "content": "hi", "tool_calls": "x"}]}, + {"x-litellm-call-id": "count-echo-utils"}, + 500, + ), + ], + ids=["counted", "counted-with-client-id", "no-messages", "failed-count-with-client-id", "crash-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 6daedc3bfbd..11c06527413 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -19885,13 +19885,13 @@ export interface paths { put?: never; /** * Token Counter - * @description Args: + * @description Returns: + * TokenCountResponse + * + * Args: * request: TokenCountRequest * 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 */ post: operations["token_counter_utils_token_counter_post"]; delete?: never;