diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index a9f95c1f9b9..bf9c119f6f2 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -206,6 +206,10 @@ class BedrockCountTokensConfig(BaseAWSLLM): encoded_body: Final = base64.b64encode(json.dumps(body_data).encode()).decode() return {"input": {"invokeModel": {"body": encoded_body}}} + @staticmethod + def count_tokens_model_id(model: str) -> str: + return get_bedrock_base_model(model).removeprefix("bedrock/") + def get_bedrock_count_tokens_endpoint( self, model: str, @@ -225,12 +229,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): Returns: Complete endpoint URL for CountTokens API """ - # Use existing LiteLLM function to get the base model ID (removes region prefix) - model_id = get_bedrock_base_model(model) - - # Remove bedrock/ prefix if present - model_id = model_id.removeprefix("bedrock/") # Remove "bedrock/" prefix - encoded_model_id: Final = self.encode_model_id(model_id=model_id) + encoded_model_id: Final = self.encode_model_id(model_id=self.count_tokens_model_id(model)) base_url, _ = self.get_runtime_endpoint( api_base=api_base, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index dd91bb6eef3..5ce6cd93a22 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3461,12 +3461,12 @@ class LiteLLM_SpendLogs_ResponseObject(LiteLLMPydanticObjectBase): class TokenCountRequest(LiteLLMPydanticObjectBase): model: str prompt: str | None = None - messages: list[dict] | None = None + messages: list[dict[Any, Any]] | None = None """ Anthropic token counting endpoint uses /messages """ - contents: list[dict] | None = None + contents: list[dict[Any, Any]] | None = None """ Google /countTokens endpoint expects contents to be a list of dicts with the following structure: """ diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 1d2045eb743..1d88e133cf9 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -18,7 +18,6 @@ from litellm.anthropic_interface.exceptions import ( AnthropicExceptionMapping, ) from litellm.integrations.custom_guardrail import ModifyResponseException -from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError from litellm.litellm_core_utils.tool_fields import first_validation_error from litellm.llms.anthropic.messages.context_management import ( AnthropicContextManagementError, @@ -49,6 +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_settings import count_source_fields from litellm.proxy.litellm_pre_call_utils import ( forwarded_anthropic_oauth_token, @@ -394,12 +394,17 @@ async def count_tokens( # Extract required fields model_name: Final = data.get("model") messages: Final = data.get("messages", []) + request_model: Final = model_name if isinstance(model_name, str) else None if not model_name: - raise HTTPException(status_code=400, detail="model parameter is required") + raise rejected_token_count( + HTTPException(status_code=400, detail="model parameter is required"), litellm_call_id, request_model + ) if not messages: - raise HTTPException(status_code=400, detail="messages parameter is required") + raise rejected_token_count( + HTTPException(status_code=400, detail="messages parameter is required"), litellm_call_id, request_model + ) auto_drop_params: Final = should_auto_drop_params_for_agentic_cli( request.headers.get("user-agent", ""), "drop_params" in data, proxy_config @@ -409,19 +414,18 @@ async def count_tokens( token_fields if isinstance(token_fields, ProxyException) else _token_count_request(token_fields) ) if isinstance(token_request, ProxyException): - return _anthropic_error_json_response(token_request, request) + return rejected_token_count( + _anthropic_error_json_response(token_request, request), litellm_call_id, request_model + ) token_response: Final = await count_request_tokens( request=token_request, call_endpoint=True, messages_format="anthropic", + litellm_call_id=litellm_call_id, client_oauth_token=forwarded_anthropic_oauth_token(request.headers), **_count_tokens_body_options(data, auto_drop_params), ) - except UnsupportedContentTypeError as e: - return _anthropic_error_json_response( - ProxyException(message=str(e), type="invalid_request_error", param="messages", code=400), request - ) except HTTPException as e: return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) except ProxyException as e: diff --git a/litellm/proxy/common_utils/token_count_log.py b/litellm/proxy/common_utils/token_count_log.py new file mode 100644 index 00000000000..04c07a29117 --- /dev/null +++ b/litellm/proxy/common_utils/token_count_log.py @@ -0,0 +1,133 @@ +import json +import re +from asyncio import CancelledError +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, TypeVar + +from fastapi import status + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.openai_error_payload import error_status_code +from litellm.types.token_count import TokenCountSource +from litellm.types.utils import TokenCountResponse + +_BARE_LOGFMT_VALUE: Final = re.compile(r"[\w.:/@+-]+", re.ASCII) +_JSON_ESCAPED_BRACES: Final = str.maketrans({"{": "\\u007b", "}": "\\u007d"}) + + +def _logfmt_value(value: str | None) -> str: + if value is None: + return "None" + if _BARE_LOGFMT_VALUE.fullmatch(value): + return value + return json.dumps(value).translate(_JSON_ESCAPED_BRACES) + + +_Counted: Final = TypeVar("_Counted") +_Rejection: Final = TypeVar("_Rejection") +_ProviderCount: Final = TypeVar("_ProviderCount", bound=Mapping[str, object]) + + +@dataclass(frozen=True, slots=True) +class _CountFields: + model_used: str | None + input_tokens: int | None + count_source: TokenCountSource | None + + +_FAILED_COUNT: Final = _CountFields(model_used=None, input_tokens=None, count_source=None) + + +@dataclass(frozen=True, slots=True) +class _CountRequest: + litellm_call_id: str + request_model: str | None + model_id: str | None + + +def deployment_model_id(deployment: Mapping[str, object] | None) -> str | None: + match deployment: + case {"model_info": {"id": str() as model_id}}: + return model_id + case _: + return None + + +def _log_token_count(request: _CountRequest, fields: _CountFields, status_code: int) -> None: + verbose_proxy_logger.info( + "token count - litellm_call_id=%s request_model=%s model_id=%s model_used=%s input_tokens=%s " + "count_source=%s status_code=%s", + _logfmt_value(request.litellm_call_id), + _logfmt_value(request.request_model), + _logfmt_value(request.model_id), + _logfmt_value(fields.model_used), + fields.input_tokens, + fields.count_source, + status_code, + extra=MappingProxyType( + { + "litellm_call_id": request.litellm_call_id, + "request_model": request.request_model, + "model_id": request.model_id, + "model_used": fields.model_used, + "input_tokens": fields.input_tokens, + "count_source": fields.count_source, + "status_code": status_code, + } + ), + ) + + +def rejected_token_count(rejection: _Rejection, litellm_call_id: str, request_model: str | None) -> _Rejection: + _log_token_count( + _CountRequest(litellm_call_id, request_model, model_id=None), + _FAILED_COUNT, + error_status_code(rejection, status.HTTP_500_INTERNAL_SERVER_ERROR), + ) + return rejection + + +async def _logged( + count: Awaitable[_Counted], request: _CountRequest, fields_of: Callable[[_Counted], _CountFields] +) -> _Counted: + try: + counted: Final = await count + except CancelledError: + raise + except BaseException as e: + _log_token_count(request, _FAILED_COUNT, error_status_code(e, status.HTTP_500_INTERNAL_SERVER_ERROR)) + raise + _log_token_count(request, fields_of(counted), status.HTTP_200_OK) + return counted + + +def _token_count_response_fields(counted: TokenCountResponse) -> _CountFields: + return _CountFields( + model_used=counted.model_used, input_tokens=counted.total_tokens, count_source=counted.count_source + ) + + +async def logged_token_count( + count: Awaitable[TokenCountResponse], litellm_call_id: str, request_model: str, model_id: str | None +) -> TokenCountResponse: + return await _logged(count, _CountRequest(litellm_call_id, request_model, model_id), _token_count_response_fields) + + +async def logged_provider_token_count( + count: Awaitable[_ProviderCount], + litellm_call_id: str, + request_model: str, + model_id: str | None, + model_used: Callable[[], str], +) -> _ProviderCount: + def provider_fields(counted: _ProviderCount) -> _CountFields: + input_tokens: Final = counted.get("input_tokens") + return _CountFields( + model_used=model_used(), + input_tokens=input_tokens if isinstance(input_tokens, int) and not isinstance(input_tokens, bool) else None, + count_source="provider", + ) + + return await _logged(count, _CountRequest(litellm_call_id, request_model, model_id), provider_fields) diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 3d1a9ab4cfe..54976857584 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -21,6 +21,8 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export 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_settings import count_source_fields from litellm.types.llms.vertex_ai import TokenCountDetailsResponse @@ -286,25 +288,35 @@ 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)) data: Final = await read_request_body(request=request) try: token_request: Final = _google_count_tokens_request(model_name, data) except ValidationError: - return _google_error_response( - request, - ProxyException( - message="Invalid Google countTokens request body", type="invalid_request_error", param=None, code=400 + return rejected_token_count( + _google_error_response( + request, + ProxyException( + message="Invalid Google countTokens request body", + type="invalid_request_error", + param=None, + code=400, + ), ), + litellm_call_id, + model_name, ) try: - token_response: Final = await internal_token_counter(request=token_request, call_endpoint=True) + token_response: Final = await internal_token_counter( + request=token_request, call_endpoint=True, litellm_call_id=litellm_call_id + ) except ProxyException as error: return _google_error_response(request, error) except HTTPException as error: return _google_error_response(request, proxy_exception_from_http_exception(error, {})) except Exception as error: - log_llm_api_exception(error, resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))) + log_llm_api_exception(error, litellm_call_id) return _google_error_response( request, ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500), diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 9eef4f0ef9d..e478096f7ad 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -87,6 +87,7 @@ from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, model_deployments_support_stream_options, open_sse_before_first_byte, + resolve_litellm_call_id, ) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -99,10 +100,16 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg safe_get_request_headers, safe_set_request_parsed_body, ) +from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.proxy.common_utils.sse_keepalive import ( wrap_passthrough_sse_bytes_with_keepalive_pings, ) +from litellm.proxy.common_utils.token_count_log import ( + deployment_model_id, + logged_provider_token_count, + rejected_token_count, +) from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key from litellm.proxy.pass_through_endpoints.injected_stream_usage import ( stream_usage_relay, @@ -1259,6 +1266,7 @@ async def handle_bedrock_count_tokens( 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() @@ -1266,40 +1274,45 @@ async def handle_bedrock_count_tokens( # Extract model from request body model: Final = _optional_str(request_body.get("model")) if not model: - raise HTTPException(status_code=400, detail={"error": "Model is required in request body"}) + raise rejected_token_count( + HTTPException(status_code=400, detail={"error": "Model is required in request body"}), + litellm_call_id, + None, + ) # Get model parameters from router litellm_params: Final = {"user_api_key_dict": user_api_key_dict} resolved_model = model # Default fallback - if llm_router: - deployments: Final = llm_router.get_model_list(model_name=model) - if deployments and len(deployments) > 0: - # Get the first matching deployment - deployment: Final = deployments[0] - model_litellm_params: Final = deployment.get("litellm_params", {}) + deployments: Final = llm_router.get_model_list(model_name=model) if llm_router else None + deployment: Final = deployments[0] if deployments else None + if deployment is not None: + model_litellm_params: Final = deployment.get("litellm_params", {}) - # Get the resolved model ID from the configuration - if "model" in model_litellm_params: - resolved_model = model_litellm_params["model"] + # Get the resolved model ID from the configuration + if "model" in model_litellm_params: + resolved_model = model_litellm_params["model"] - # Copy all litellm_params - BaseAWSLLM will handle AWS credential discovery - for key, value in model_litellm_params.items(): - if key != "user_api_key_dict": # Don't overwrite user_api_key_dict - litellm_params[key] = value + # Copy all litellm_params - BaseAWSLLM will handle AWS credential discovery + for key, value in model_litellm_params.items(): + if key != "user_api_key_dict": # Don't overwrite user_api_key_dict + litellm_params[key] = value verbose_proxy_logger.debug("Count tokens litellm_params: %s", litellm_params) verbose_proxy_logger.debug("Resolved model: %s", resolved_model) - # Handle the count tokens request - result: Final = await handler.handle_count_tokens_request( - request_data=request_body, - litellm_params=litellm_params, - resolved_model=resolved_model, + return await logged_provider_token_count( + handler.handle_count_tokens_request( + request_data=request_body, + litellm_params=litellm_params, + resolved_model=resolved_model, + ), + litellm_call_id=litellm_call_id, + request_model=model, + model_id=deployment_model_id(deployment), + model_used=partial(handler.count_tokens_model_id, resolved_model), ) - return result - except BedrockError as e: # Convert BedrockError to HTTPException for FastAPI verbose_proxy_logger.error("BedrockError in handle_bedrock_count_tokens: %s", e) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4e7c4601803..7fcb77e2bbf 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -100,7 +100,7 @@ from litellm.litellm_core_utils.litellm_logging import ( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError, offload_token_count from litellm.llms.bedrock.count_tokens.kimi_k3 import ( HUGGING_FACE_SOURCE, KimiK3Encodings, @@ -552,6 +552,11 @@ from litellm.proxy.common_utils.timezone_utils import ( get_budget_reset_settings, get_budget_reset_time, ) +from litellm.proxy.common_utils.token_count_log import ( + deployment_model_id, + logged_token_count, + rejected_token_count, +) from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, end_user_cache_key, @@ -15253,16 +15258,21 @@ def _system_message(system: object) -> ChatCompletionSystemMessage | None: dependencies=[Depends(user_api_key_auth)], response_model=TokenCountResponse, ) -async def token_counter(request: TokenCountRequest, call_endpoint: bool = False): +async def token_counter( + request: TokenCountRequest, + call_endpoint: bool = False, + litellm_call_id: Annotated[str | None, Header(alias=LITELLM_CALL_ID_HEADER)] = None, +): """ 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 """ - return await count_request_tokens(request=request, call_endpoint=call_endpoint) + return await count_request_tokens(request=request, call_endpoint=call_endpoint, litellm_call_id=litellm_call_id) async def count_request_tokens( @@ -15273,7 +15283,69 @@ async def count_request_tokens( context_management: JsonValue = None, request_litellm_params: Mapping[str, JsonValue] = MappingProxyType({}), responses_request: Mapping[str, object] | None = None, + litellm_call_id: str | None = None, client_oauth_token: str | None = None, +) -> TokenCountResponse: + call_id: Final = resolve_litellm_call_id(litellm_call_id) + if request.prompt is None and request.messages is None and request.contents is None: + raise rejected_token_count( + HTTPException(status_code=400, detail="prompt or messages or contents must be provided"), + call_id, + request.model, + ) + routed: Final = await _routed_count_deployment(request.model, responses_request) + count: Final = _count_request_tokens( + request=request, + routed=routed, + call_endpoint=call_endpoint, + messages_format=messages_format, + output_format=output_format, + context_management=context_management, + request_litellm_params=request_litellm_params, + responses_request=responses_request, + client_oauth_token=client_oauth_token, + ) + return await logged_token_count( + _with_anthropic_content_errors(count) if messages_format == "anthropic" else count, + litellm_call_id=call_id, + request_model=request.model, + model_id=deployment_model_id(routed), + ) + + +async def _with_anthropic_content_errors(count: Awaitable[TokenCountResponse]) -> TokenCountResponse: + try: + return await count + except UnsupportedContentTypeError as e: + raise ProxyException(message=str(e), type="invalid_request_error", param="messages", code=400) from e + + +async def _routed_count_deployment( + model: str, responses_request: Mapping[str, object] | None +) -> Mapping[str, object] | None: + if llm_router is None: + return None + try: + return await llm_router.async_get_available_deployment( + model=model, request_kwargs=dict(_responses_routing_kwargs(responses_request)) + ) + except Exception: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.token_counter(): Exception occured while getting deployment" + ) + return None + + +async def _count_request_tokens( + request: TokenCountRequest, + routed: Mapping[str, object] | None, + call_endpoint: bool, + messages_format: TokenCountMessagesFormat, + output_format: JsonValue, + context_management: JsonValue, + request_litellm_params: Mapping[str, JsonValue], + responses_request: Mapping[str, object] | None, + client_oauth_token: str | None, ) -> TokenCountResponse: prompt: Final = request.prompt messages: Final = request.messages @@ -15281,26 +15353,9 @@ async def count_request_tokens( tools: Final = request.tools system: Final = request.system - ######################################################### - # Validate request - ######################################################### - if prompt is None and messages is None and contents is None: - raise HTTPException(status_code=400, detail="prompt or messages or contents must be provided") - - deployment: dict[str, Any] | None = None + deployment: dict[str, Any] | None = None if routed is None else dict(routed) litellm_model_name = None model_info: ModelMapInfo | None = None - if llm_router is not None: - # get 1 deployment corresponding to the model - try: - deployment = await llm_router.async_get_available_deployment( - model=request.model, - request_kwargs=dict(_responses_routing_kwargs(responses_request)), - ) - except Exception: - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.token_counter(): Exception occured while getting deployment" - ) if deployment is not None: litellm_model_name = deployment.get("litellm_params", {}).get("model") model_info = deployment.get("model_info", {}) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 97bbf893e3e..4f8b1697f7e 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -56,6 +56,7 @@ from litellm.proxy.common_request_processing import ( cost_headers, create_response, resolve_counted_tool_fields, + resolve_litellm_call_id, ) from litellm.proxy.common_utils.error_body_call_id import error_body_call_id, with_call_id from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports @@ -69,6 +70,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, openai_error_type_for_status, ) +from litellm.proxy.common_utils.token_count_log import 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 @@ -1407,7 +1409,12 @@ class _TokenCountPayload(TypedDict): class _TokenCounter(Protocol): def __call__( - self, *, request: TokenCountRequest, call_endpoint: bool, responses_request: Mapping[str, object] | None + self, + *, + request: TokenCountRequest, + call_endpoint: bool, + responses_request: Mapping[str, object] | None, + litellm_call_id: str | None, ) -> Awaitable[TokenCountResponse]: ... @@ -1648,18 +1655,23 @@ 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)) 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 _missing_responses_param_response("model") + return rejected_token_count(_missing_responses_param_response("model"), litellm_call_id, None) if input_value is None: - return _missing_responses_param_response("input") + return rejected_token_count(_missing_responses_param_response("input"), litellm_call_id, model_name) if isinstance(input_value, (str, list)) and not input_value: - return _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", + 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", + ), + litellm_call_id, + model_name, ) counted_fields: Final = { @@ -1667,7 +1679,11 @@ 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 _responses_invalid_request_response(message=token_fields.message, param=token_fields.param, code=None) + return rejected_token_count( + _responses_invalid_request_response(message=token_fields.message, param=token_fields.param, code=None), + litellm_call_id, + model_name, + ) try: payload: Final[_TokenCountPayload] = { @@ -1680,27 +1696,44 @@ async def responses_input_tokens( } token_request: Final = TokenCountRequest.model_validate(payload) except ValidationError as e: - return _responses_invalid_request_response( - message=first_validation_error(e), param=str(e.errors(include_url=False)[0]["loc"][0]), code=None + 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 + ), + litellm_call_id, + model_name, ) except Exception as e: - return _responses_invalid_request_response( - message=f"Invalid request for token counting: {e}", param=None, code=None + return rejected_token_count( + _responses_invalid_request_response( + message=f"Invalid request for token counting: {e}", param=None, code=None + ), + litellm_call_id, + model_name, ) request_fields: Final = _PROXY_SERVER_REQUEST_ADAPTER.validate_python(data) - responses_request: Final = await _with_managed_file_id_mapping( - _with_authorized_previous_response( - {**{name: value for name, value in request_fields.items() if name not in counted_fields}, **token_fields}, - response_id_security, + try: + responses_request: Final = await _with_managed_file_id_mapping( + _with_authorized_previous_response( + { + **{name: value for name, value in request_fields.items() if name not in counted_fields}, + **token_fields, + }, + response_id_security, + user_api_key_dict, + ), + managed_files, user_api_key_dict, - ), - managed_files, - user_api_key_dict, - ) + ) + except ProxyException as e: + raise rejected_token_count(e, litellm_call_id, model_name) try: token_response: Final = await token_counter( - request=token_request, call_endpoint=True, responses_request=responses_request + request=token_request, + call_endpoint=True, + responses_request=responses_request, + litellm_call_id=litellm_call_id, ) except ProxyException as e: if e.type != _PROVIDER_COUNT_ERROR_TYPE: diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index c7bb05e5d30..037b0a6c29a 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -5,9 +5,11 @@ Test for anthropic_endpoints/endpoints.py, focusing on handling dictionary objec import json import logging import unittest +from collections.abc import Callable, Mapping from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi.testclient import TestClient @@ -615,6 +617,184 @@ async def test_count_tokens_drops_malformed_tools_the_counted_deployment_drops_e assert counter.await_args.kwargs["request"].tools is None +def _count_on_claude( + monkeypatch: pytest.MonkeyPatch, peer: Callable[[httpx.Request], httpx.Response], headers: Mapping[str, str] +) -> httpx.Response: + from fastapi import FastAPI + + import litellm.proxy.proxy_server as proxy_server + from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.anthropic_endpoints.endpoints import router + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + claude: Final = {"model": "anthropic/claude-zeta-9", "api_key": "test-key", "api_base": "https://count-tokens.test"} + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "count-log-claude", + "litellm_params": {**claude, "order": 2}, + "model_info": {"id": "count-log-fallback-deployment"}, + }, + { + "model_name": "count-log-claude", + "litellm_params": {**claude, "order": 1}, + "model_info": {"id": "count-log-routed-deployment"}, + }, + ] + ), + ) + monkeypatch.setattr( + anthropic_count_tokens_handler, + "get_async_httpx_client", + lambda *args, **kwargs: AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ) + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth() + return TestClient(app).post( + "/v1/messages/count_tokens", + json={"model": "count-log-claude", "messages": [{"role": "user", "content": "Hello"}]}, + headers=dict(headers), + ) + + +def test_count_tokens_logs_one_line_with_the_callers_call_id_model_and_count( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _count_on_claude( + monkeypatch, + lambda request: httpx.Response(200, json={"input_tokens": 77}, request=request), + {"x-litellm-call-id": "count-log-msg-claude"}, + ) + + assert response.json() == {"input_tokens": 77, "litellm_count_source": "provider"}, response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-msg-claude", + "request_model": "count-log-claude", + "model_id": "count-log-routed-deployment", + "model_used": "claude-zeta-9", + "input_tokens": 77, + "count_source": "provider", + "status_code": 200, + }, + ) + + +def test_count_tokens_provider_rejection_logs_its_status_and_answers_the_same_error( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + rejection: Final = { + "type": "error", + "error": {"type": "invalid_request_error", "message": "messages.0.content.0: invalid base64 data"}, + } + + response: Final = _count_on_claude( + monkeypatch, + lambda request: httpx.Response(400, json=rejection, request=request), + {"x-litellm-call-id": "count-log-msg-error"}, + ) + + assert response.status_code == 400, response.text + assert response.json()["error"] == rejection["error"] + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-msg-error", + "request_model": "count-log-claude", + "model_id": "count-log-routed-deployment", + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + +def _count_on_local_model(monkeypatch: pytest.MonkeyPatch, body: Mapping[str, object]) -> httpx.Response: + from fastapi import FastAPI + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.anthropic_endpoints.endpoints import router + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "count-log-local", + "litellm_params": {"model": "hosted_vllm/count-log-local", "api_base": "https://local.test/v1"}, + "model_info": {"id": "count-log-local-deployment"}, + } + ] + ), + ) + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth() + return TestClient(app).post( + "/v1/messages/count_tokens", json=dict(body), headers={"x-litellm-call-id": "count-log-local"} + ) + + +def test_count_tokens_unsupported_content_logs_the_400_it_answers( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _count_on_local_model( + monkeypatch, + {"model": "count-log-local", "messages": [{"role": "user", "content": [{"type": "future_block"}]}]}, + ) + + assert (response.status_code, response.json()["error"]["type"]) == (400, "invalid_request_error"), response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-local", + "request_model": "count-log-local", + "model_id": "count-log-local-deployment", + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + +@pytest.mark.parametrize( + ("body", "request_model"), + [ + ({"messages": [{"role": "user", "content": "Hello"}]}, None), + ({"model": "count-log-local", "messages": []}, "count-log-local"), + ], +) +def test_count_tokens_rejected_before_counting_logs_the_400_it_answers( + monkeypatch: pytest.MonkeyPatch, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], + body: Mapping[str, object], + request_model: str | None, +) -> None: + response: Final = _count_on_local_model(monkeypatch, body) + + assert response.status_code == 400, response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-local", + "request_model": request_model, + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + 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 new file mode 100644 index 00000000000..94b976d88cf --- /dev/null +++ b/tests/unit/proxy/common_utils/test_token_count_log.py @@ -0,0 +1,45 @@ +import asyncio +from collections.abc import Callable + +import pytest + +from litellm.proxy.common_utils.token_count_log import logged_token_count +from litellm.types.utils import TokenCountResponse + + +class _TokenizerPanic(BaseException): + pass + + +async def _count_ending_in(error: BaseException) -> TokenCountResponse: + raise error + + +@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], ...]], +) -> None: + with pytest.raises(_TokenizerPanic): + await logged_token_count(_count_ending_in(_TokenizerPanic()), "count-log-panic", "count-log-gpt", None) + + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-panic", + "request_model": "count-log-gpt", + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 500, + }, + ) + + +@pytest.mark.asyncio +async def test_a_cancelled_count_writes_no_line( + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + with pytest.raises(asyncio.CancelledError): + await logged_token_count(_count_ending_in(asyncio.CancelledError()), "count-log-cancelled", "count-log-gpt", None) + + assert token_count_log_lines() == () diff --git a/tests/unit/proxy/conftest.py b/tests/unit/proxy/conftest.py index 58c9a5e4d76..e13eddef43b 100644 --- a/tests/unit/proxy/conftest.py +++ b/tests/unit/proxy/conftest.py @@ -3,6 +3,7 @@ import asyncio import copy import inspect +import logging import os import tempfile import warnings @@ -17,6 +18,7 @@ from prisma.errors import ClientNotConnectedError import litellm import litellm.proxy.proxy_server +from litellm._logging import verbose_proxy_logger from litellm._service_logger import ServiceTypes from litellm.integrations.otel.model.payloads import ServiceSpanData from litellm.integrations.otel.model.spans import service_span_name @@ -424,6 +426,31 @@ def fresh_agent_read_through(monkeypatch): return read_through +TOKEN_COUNT_LOG_FIELDS: Final = ( + "litellm_call_id", + "request_model", + "model_id", + "model_used", + "input_tokens", + "count_source", + "status_code", +) + + +@pytest.fixture +def token_count_log_lines( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> Iterator[Callable[[], tuple[dict[str, object], ...]]]: + """The structured fields of every INFO token count line the proxy logger writes while the test runs.""" + monkeypatch.setattr(verbose_proxy_logger, "propagate", True) + with caplog.at_level(logging.INFO, logger=verbose_proxy_logger.name): + yield lambda: tuple( + {field: record.__dict__[field] for field in TOKEN_COUNT_LOG_FIELDS} + for record in caplog.records + if record.levelno == logging.INFO and "count_source" in record.__dict__ + ) + + @pytest.fixture def postgres_span_names() -> Iterator[Callable[[], Awaitable[tuple[str, ...]]]]: """The ``postgres.{verb} {table}`` names OTel would render for every DB service event diff --git a/tests/unit/proxy/google_endpoints/test_endpoints.py b/tests/unit/proxy/google_endpoints/test_endpoints.py index 17af52b5579..80e45edf897 100644 --- a/tests/unit/proxy/google_endpoints/test_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_endpoints.py @@ -94,6 +94,7 @@ async def test_proxy_gemini_to_openai_like_model_token_counting(): request=Request( scope={ "type": "http", + "headers": [], "parsed_body": ( ["contents"], {"contents": [{"parts": [{"text": "Hello, how are you?"}]}]}, 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 2777d18a399..4867ca9ab63 100644 --- a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py @@ -3,9 +3,11 @@ Test to verify the Google GenAI proxy API endpoints """ +from collections.abc import Callable from typing import Final from unittest.mock import AsyncMock, patch +import httpx import pytest from litellm.types.utils import TokenCountResponse @@ -240,3 +242,79 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs() assert body["promptTokensDetails"] == [] assert body["litellm_count_source"] == "token_count_model" assert body["litellm_count_model"] == "us.anthropic.claude-opus-5-5" + + +@pytest.mark.parametrize("path", ["/v1beta/models/count-log-claude:countTokens", "/models/count-log-claude:countTokens"]) +def test_google_count_tokens_logs_one_line_with_the_callers_call_id_model_and_count( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], path: str +) -> None: + import litellm.proxy.proxy_server as proxy_server + from litellm import Router + from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + claude: Final = {"model": "anthropic/claude-zeta-9", "api_key": "test-key", "api_base": "https://count-tokens.test"} + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "count-log-claude", + "litellm_params": claude, + "model_info": {"id": "count-log-gemini-deployment"}, + } + ] + ), + ) + monkeypatch.setattr( + anthropic_count_tokens_handler, + "get_async_httpx_client", + lambda *args, **kwargs: AsyncHTTPHandler( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, json={"input_tokens": 77}, request=request) + ) + ), + ) + + response: Final = _build_test_client().post( + path, + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + headers={"x-litellm-call-id": "count-log-gemini"}, + ) + + assert (response.status_code, response.json()["totalTokens"]) == (200, 77), response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-gemini", + "request_model": "count-log-claude", + "model_id": "count-log-gemini-deployment", + "model_used": "claude-zeta-9", + "input_tokens": 77, + "count_source": "provider", + "status_code": 200, + }, + ) + + +def test_google_count_tokens_rejected_body_logs_the_400_it_answers( + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + response: Final = _build_test_client().post( + "/v1beta/models/count-log-claude:countTokens", + json={"contents": "not a list"}, + headers={"x-litellm-call-id": "count-log-gemini-bad-body"}, + ) + + assert response.status_code == 400, response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-gemini-bad-body", + "request_model": "count-log-claude", + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) 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 60043daa402..43fe5d4e786 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 @@ -2142,7 +2142,6 @@ class TestBedrockPassthroughAwsErrors: assert exc_info.value.aws_error_type == "ValidationException" assert (exc_info.value.error.message, exc_info.value.error.code) == (message, status_code) - @pytest.mark.asyncio async def test_bedrock_request_check_is_raised_with_the_aws_error_type_it_carries( self, monkeypatch: pytest.MonkeyPatch @@ -2183,6 +2182,7 @@ class TestBedrockPassthroughAwsErrors: await self._answer(error) assert (exc_info.value.code, exc_info.value.headers.get("x-model-group")) == ("403", "modèle") + @pytest.mark.asyncio @pytest.mark.parametrize( ("error", "message", "code"), @@ -7250,6 +7250,189 @@ async def test_bedrock_count_tokens_error_forwards_provider_headers(): assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-count-tokens-500" +_BEDROCK_COUNT_ALIAS: Final = "count-log-bedrock" +_BEDROCK_COUNT_MODEL: Final = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" + + +async def _bedrock_native_count( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, aws_reply: httpx.Response +) -> dict[str, object]: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_bedrock_count_tokens, + ) + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + litellm.Router( + model_list=[ + { + "model_name": _BEDROCK_COUNT_ALIAS, + "litellm_params": { + "model": _BEDROCK_COUNT_MODEL, + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + }, + "model_info": {"id": "count-log-bedrock-deployment"}, + } + ] + ), + ) + respx_mock.post(host="bedrock-runtime.us-east-1.amazonaws.com", path__regex=r"/count-tokens$").mock( + return_value=aws_reply + ) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": f"/bedrock/model/{_BEDROCK_COUNT_ALIAS}/count-tokens", + "headers": [(b"x-litellm-call-id", b"bedrock-native-count")], + } + ) + return 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={"model": _BEDROCK_COUNT_ALIAS, "messages": [{"role": "user", "content": "hi"}]}, + ) + + +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_logs_the_count_against_the_router_model( + monkeypatch: pytest.MonkeyPatch, + httpx_transport: None, + 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}) + ) + + assert result == {"input_tokens": 12} + 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() == ( + { + "litellm_call_id": "bedrock-native-count", + "request_model": _BEDROCK_COUNT_ALIAS, + "model_id": "count-log-bedrock-deployment", + "model_used": "anthropic.claude-haiku-4-5-20251001-v1:0", + "input_tokens": 12, + "count_source": "provider", + "status_code": 200, + }, + ) + + +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_logs_the_provider_status_when_the_count_fails( + monkeypatch: pytest.MonkeyPatch, + httpx_transport: None, + respx_mock: respx.MockRouter, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> 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) + + assert exc_info.value.status_code == 400 + assert token_count_log_lines() == ( + { + "litellm_call_id": "bedrock-native-count", + "request_model": _BEDROCK_COUNT_ALIAS, + "model_id": "count-log-bedrock-deployment", + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_without_a_model_logs_the_400_it_answers( + 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": [(b"x-litellm-call-id", b"bedrock-native-no-model")], + } + ) + 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"}]}, + ) + + assert exc_info.value.status_code == 400 + assert token_count_log_lines() == ( + { + "litellm_call_id": "bedrock-native-no-model", + "request_model": None, + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + + +@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, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_bedrock_count_tokens, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/bedrock/model/us/count-tokens", + "headers": [(b"x-litellm-call-id", b"bedrock-native-region-only")], + } + ) + with pytest.raises(HTTPException) as exc_info: + await handle_bedrock_count_tokens( + endpoint="model/us/count-tokens", + request=request, + fastapi_response=Response(), + user_api_key_dict=MagicMock(), + request_body={"model": "us"}, + ) + + assert exc_info.value.detail == { + "error": "CountTokens processing error: Request must contain content to count tokens" + } + assert token_count_log_lines() == ( + { + "litellm_call_id": "bedrock-native-region-only", + "request_model": "us", + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": exc_info.value.status_code, + }, + ) + class _AzureGroupRouter: def __init__(self, captured: list[dict]) -> None: self.captured = captured diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 985aee061bc..7760b6b7fcc 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -7,7 +7,7 @@ import contextlib import dataclasses import json import unittest -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import Any, Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, call, patch @@ -19,6 +19,7 @@ from fastapi.testclient import TestClient from httpx import Response import litellm +from litellm.constants import EMPTY_MAPPING from litellm.proxy._types import ProxyException from litellm.proxy.proxy_server import app from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse @@ -4275,6 +4276,7 @@ class TestResponsesInputTokens: llm_router: litellm.Router | None = None, response_id_security: object | None = None, managed_files: object | None = None, + headers: Mapping[str, str] = EMPTY_MAPPING, ) -> tuple[Response, AsyncMock]: from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -4307,7 +4309,7 @@ class TestResponsesInputTokens: app.dependency_overrides[_proxy_managed_files_hook] = lambda: managed_files try: client = TestClient(app) - response = client.post(path, json=body, headers={"Authorization": "Bearer sk-9876"}) + response = client.post(path, json=body, headers={"Authorization": "Bearer sk-9876", **headers}) return response, token_counter_mock finally: app.dependency_overrides.pop(user_api_key_auth, None) @@ -4552,6 +4554,84 @@ class TestResponsesInputTokens: assert (response.status_code, response.json().get("input_tokens")) == (200, 7), response.text assert sent == [{"model": "gpt-4o-mini", "input": "hi"}] + @pytest.mark.parametrize( + "path", ["/v1/responses/input_tokens", "/responses/input_tokens", "/openai/v1/responses/input_tokens"] + ) + def test_count_logs_one_line_with_the_callers_call_id_model_and_count( + self, + monkeypatch: pytest.MonkeyPatch, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], + path: str, + ) -> None: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.llms.openai.responses.count_tokens import handler as openai_count_tokens_http + from litellm.proxy.proxy_server import count_request_tokens + + gpt: Final = {"model": "openai/gpt-4o-mini", "api_key": "test-key", "api_base": "https://count.test/v1"} + router: Final = litellm.Router( + model_list=[ + { + "model_name": "count-log-gpt", + "litellm_params": gpt, + "model_info": {"id": "count-log-responses-deployment"}, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + monkeypatch.setattr( + openai_count_tokens_http, + "get_async_httpx_client", + lambda *args, **kwargs: AsyncHTTPHandler( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 200, json={"object": "response.input_tokens", "input_tokens": 7}, request=request + ) + ) + ), + ) + + response, _ = self._post_input_tokens( + {"model": "count-log-gpt", "input": "hi"}, + path=path, + counter=AsyncMock(side_effect=count_request_tokens), + llm_router=router, + headers={"x-litellm-call-id": "count-log-responses"}, + ) + + assert (response.status_code, response.json().get("input_tokens")) == (200, 7), response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-responses", + "request_model": "count-log-gpt", + "model_id": "count-log-responses-deployment", + "model_used": "gpt-4o-mini", + "input_tokens": 7, + "count_source": "provider", + "status_code": 200, + }, + ) + + def test_count_rejected_before_counting_logs_the_400_it_answers( + self, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] + ) -> None: + response, counter = self._post_input_tokens( + {"model": "count-log-gpt"}, headers={"x-litellm-call-id": "count-log-no-input"} + ) + + assert response.status_code == 400, response.text + counter.assert_not_awaited() + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-no-input", + "request_model": "count-log-gpt", + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + def test_a_managed_file_the_key_cannot_read_answers_403_without_a_count(self) -> None: from litellm.integrations.custom_logger import CustomLogger diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index ebe3c7407d8..5484bc38317 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -10,6 +10,7 @@ import json import logging import threading import traceback +import uuid from collections.abc import Callable, Coroutine, Mapping from functools import partial from pathlib import Path @@ -40,7 +41,7 @@ from starlette.types import Message import litellm from litellm import Router -from litellm._logging import verbose_proxy_logger +from litellm._logging import JsonFormatter, verbose_proxy_logger from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError @@ -364,6 +365,7 @@ async def test_anthropic_messages_count_tokens_endpoint(): output_format=None, context_management=None, request_litellm_params=None, + litellm_call_id=None, client_oauth_token=None, ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" @@ -445,6 +447,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): output_format=None, context_management=None, request_litellm_params=None, + litellm_call_id=None, client_oauth_token=None, ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" @@ -1264,6 +1267,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): output_format=None, context_management=None, request_litellm_params=None, + litellm_call_id=None, client_oauth_token=None, ): raise ProxyException( @@ -1322,6 +1326,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): output_format=None, context_management=None, request_litellm_params=None, + litellm_call_id=None, client_oauth_token=None, ): raise ProxyException( @@ -1380,6 +1385,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): output_format=None, context_management=None, request_litellm_params=None, + litellm_call_id=None, client_oauth_token=None, ): raise ProxyException( @@ -3897,3 +3903,149 @@ async def test_responses_count_on_claude_counts_the_turns_previous_response_id_r ] assert [sorted(message) for message in messages if isinstance(message, dict)] == [["content", "role"]] * 3 spend_logs.assert_awaited_once_with("chatcmpl-stored", ("chatcmpl-stored", "resp_client_form")) + + +_COUNTED_HELLO: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"model": "count-log-gpt", "messages": [{"role": "user", "content": "Hello, how many tokens is this?"}]} +) + + +def _utils_token_counter( + monkeypatch: pytest.MonkeyPatch, body: Mapping[str, JsonValue], headers: Mapping[str, str] +) -> httpx.Response: + from fastapi.testclient import TestClient + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import app + + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "count-log-gpt", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, + "model_info": {"id": "count-log-utils-deployment"}, + } + ] + ), + ) + 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)) + + +def test_utils_token_counter_logs_one_line_with_the_callers_call_id_model_and_count( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _utils_token_counter(monkeypatch, _COUNTED_HELLO, {"x-litellm-call-id": "count-log-utils"}) + + assert response.status_code == 200, response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-utils", + "request_model": "count-log-gpt", + "model_id": "count-log-utils-deployment", + "model_used": "gpt-4o-mini", + "input_tokens": response.json()["total_tokens"], + "count_source": "local_estimate", + "status_code": 200, + }, + ) + + +def test_utils_token_counter_logs_a_generated_call_id_when_the_client_sends_none( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _utils_token_counter(monkeypatch, _COUNTED_HELLO, {}) + + assert response.status_code == 200, response.text + (line,) = token_count_log_lines() + assert str(uuid.UUID(str(line["litellm_call_id"]))) == line["litellm_call_id"] + + +def test_utils_token_counter_without_messages_logs_the_400_it_answers( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _utils_token_counter( + monkeypatch, {"model": "count-log-gpt"}, {"x-litellm-call-id": "count-log-utils-empty"} + ) + + assert response.status_code == 400, response.text + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-utils-empty", + "request_model": "count-log-gpt", + "model_id": None, + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + +def test_utils_token_counter_failed_count_logs_its_status_in_place_of_a_count( + monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]] +) -> None: + response: Final = _utils_token_counter( + monkeypatch, {**_COUNTED_HELLO, "prompt": "and a prompt"}, {"x-litellm-call-id": "count-log-utils-error"} + ) + + assert response.status_code == 400, response.text + assert response.json() == { + "error": { + "message": "prompt cannot be sent with messages or contents. Set drop_params to count them without the prompt", + "type": "invalid_request_error", + "param": "prompt", + "code": "400", + } + } + assert token_count_log_lines() == ( + { + "litellm_call_id": "count-log-utils-error", + "request_model": "count-log-gpt", + "model_id": "count-log-utils-deployment", + "model_used": None, + "input_tokens": None, + "count_source": None, + "status_code": 400, + }, + ) + + +def test_utils_token_counter_json_line_keeps_its_own_fields_when_the_call_id_is_a_dict( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + forged_call_id: Final = "{'component':'forged-comp','logger':'forged.py:1','injected_key':'injected'}" + + _utils_token_counter(monkeypatch, _COUNTED_HELLO, {"x-litellm-call-id": forged_call_id}) + + (line,) = token_count_log_lines() + (record,) = (record for record in caplog.records if "count_source" in record.__dict__) + rendered: Final = json.loads(JsonFormatter().format(record)) + assert line["litellm_call_id"] == forged_call_id + assert {key: rendered.get(key) for key in ("litellm_call_id", "component", "logger", "injected_key")} == { + "litellm_call_id": forged_call_id, + "component": verbose_proxy_logger.name, + "logger": f"{record.filename}:{record.lineno}", + "injected_key": None, + } + + +def test_utils_token_counter_line_stays_one_line_when_the_model_carries_a_newline( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], +) -> None: + forged_model: Final = "count-log-gpt\ntoken count - litellm_call_id=forged input_tokens=1" + + _utils_token_counter(monkeypatch, {**_COUNTED_HELLO, "model": forged_model}, {"x-litellm-call-id": "count log"}) + + (line,) = token_count_log_lines() + (message,) = (record.getMessage() for record in caplog.records if "count_source" in record.__dict__) + assert line["request_model"] == forged_model + assert "\n" not in message + assert message.startswith(f'token count - litellm_call_id="count log" request_model={json.dumps(forged_model)} ') diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 76e3c95fb75..f9117afdf3a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -19058,6 +19058,7 @@ export interface paths { * @description 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 @@ -76773,7 +76774,9 @@ export interface operations { query?: { call_endpoint?: boolean; }; - header?: never; + header?: { + "x-litellm-call-id"?: string | null; + }; path?: never; cookie?: never; };