scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

feat(proxy): log one line per token count call (#47657)

Browse source 
* feat(proxy): log one line per token count call

Each count call on /v1/messages/count_tokens, /utils/token_counter, /v1/responses/input_tokens and the Gemini countTokens routes now logs one INFO line with the litellm call id, the requested model, the model the count ran against, the count, the count source and the status code, so it joins the inference call that sends the same x-litellm-call-id

* test(google_endpoints): give the count test's fake request a headers list

* fix(proxy): keep token count log lines one per call and trusted in JSON logs, and log Bedrock native counts

Quote log values that carry whitespace, quotes, equals, or control characters so a model name can't split or forge a line. Let JsonFormatter record extras win over a dict parsed from the message. Log the Bedrock native count-tokens route through the same line.

* fix(proxy): escape braces in token count log values, log the Bedrock model AWS counts, and skip the line when there is nothing to count

* fix(proxy): log the deployment id the count routed to on every token count line

* fix(proxy): log count requests rejected before counting with the status they answer

A count route that rejects a request before the shared count function runs now logs the same token count line with its 400. An unsupported content block on the Messages count route becomes its 400 inside the logged wrapper, so its line says 400 instead of 500. A count that ends in a BaseException other than cancellation logs a 500.

* fix(proxy): compute the Bedrock counted model id only after a successful count

* revert(logging): keep the JSON log formatter's field order as on main

The count line already escapes braces in its values, so the formatter never parses a dict from it. Changing the order for every logger is a separate decision.
This commit is contained in:
Mateo Wang 2026-10-28 11:23:52 -07:00 • committed by GitHub
parent 6e33d7e84b
commit ea6dbd2109
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 1088 additions and 90 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/llms/bedrock/count_tokens/transformation.py
View file

@ -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,

4
litellm/proxy/_types.py
View file

@ -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:
"""

20
litellm/proxy/anthropic_endpoints/endpoints.py
View file

@ -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:

133
litellm/proxy/common_utils/token_count_log.py Normal file
View file

@ -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)

24
litellm/proxy/google_endpoints/endpoints.py
View file

@ -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),

55
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -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)

97
litellm/proxy/proxy_server.py
View file

@ -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", {})

75
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -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": <count>}`
"""
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:

180
tests/unit/proxy/anthropic_endpoints/test_endpoints.py
View file

@ -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"""

45
tests/unit/proxy/common_utils/test_token_count_log.py Normal file
View file

@ -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() == ()

27
tests/unit/proxy/conftest.py
View file

@ -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

1
tests/unit/proxy/google_endpoints/test_endpoints.py
View file

@ -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?"}]}]},

78
tests/unit/proxy/google_endpoints/test_google_api_endpoints.py
View file

@ -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,
},
)

185
tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
View file

@ -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

84
tests/unit/proxy/response_api_endpoints/test_endpoints.py
View file

@ -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

154
tests/unit/proxy/test_proxy_token_counter.py
View file

@ -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)} ')

5
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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;
};