scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

fix(proxy): answer x-litellm-call-id on every token count response

Browse source 
The token count log line names a litellm_call_id, but count responses
never carried it, so a client that sent no id of its own could not find
its count's log line. Every count route (Anthropic Messages, Responses
input_tokens, Gemini countTokens, Bedrock count-tokens, and
/utils/token_counter) now answers the id on success and on errors, and
a crash answers the JSON 500 the Anthropic and Gemini routes already used.
This commit is contained in:
mateo-berri 2026-10-29 08:26:20 -07:00
parent 67ae3fb527
commit d13b74161a
15 changed files with 435 additions and 106 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -48,7 +48,7 @@ from litellm.proxy.common_utils.openai_error_payload import (
openai_error_type,
with_litellm_call_id,
)
from litellm.proxy.common_utils.token_count_log import rejected_token_count
from litellm.proxy.common_utils.token_count_log import answered_with_litellm_call_id, rejected_token_count
from litellm.proxy.common_utils.token_count_settings import count_source_fields
from litellm.proxy.litellm_pre_call_utils import (
forwarded_anthropic_oauth_token,
@ -362,6 +362,7 @@ def _count_tokens_body_options(data: Mapping[str, JsonValue], auto_drop_params:
)
async def count_tokens(
request: Request,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # Used for auth
):
"""
@ -384,9 +385,17 @@ async def count_tokens(
Returns: {"input_tokens": <number>}
"""
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER))
return await answered_with_litellm_call_id(
_count_tokens(request, user_api_key_dict, litellm_call_id), litellm_call_id, fastapi_response
)
async def _count_tokens(
request: Request, user_api_key_dict: UserAPIKeyAuth, litellm_call_id: str
) -> JSONResponse | Mapping[str, object]:
from litellm.proxy.proxy_server import count_request_tokens, llm_router, proxy_config
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
try:
request_data: Final = await read_request_body(request=request)
data: Final[dict] = {**request_data}
@ -415,7 +424,9 @@ async def count_tokens(
)
if isinstance(token_request, ProxyException):
return rejected_token_count(
_anthropic_error_json_response(token_request, request), litellm_call_id, request_model
_anthropic_error_json_response(with_litellm_call_id(token_request, litellm_call_id), request),
litellm_call_id,
request_model,
)
token_response: Final = await count_request_tokens(
@ -427,12 +438,16 @@ async def count_tokens(
**_count_tokens_body_options(data, auto_drop_params),
)
except HTTPException as e:
return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request)
return _anthropic_error_json_response(
with_litellm_call_id(proxy_exception_from_http_exception(e, {}), litellm_call_id), request
)
except ProxyException as e:
return _anthropic_error_json_response(e, request)
return _anthropic_error_json_response(with_litellm_call_id(e, litellm_call_id), request)
except Exception as e:
log_llm_api_exception(e, litellm_call_id)
return _anthropic_error_json_response(proxy_exception_from_route_error(e), request)
return _anthropic_error_json_response(
with_litellm_call_id(proxy_exception_from_route_error(e), litellm_call_id), request
)
else:
return {"input_tokens": token_response.total_tokens, **count_source_fields(token_response)}

35
litellm/proxy/common_utils/token_count_log.py
View file

@ -6,10 +6,17 @@ from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, TypeVar
from fastapi import status
from fastapi import HTTPException, Response, status
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_utils.openai_error_payload import error_status_code
from litellm.proxy._types import ProxyException
from litellm.proxy.common_request_processing import log_llm_api_exception, proxy_exception_from_route_error
from litellm.proxy.common_utils.openai_error_payload import (
LITELLM_CALL_ID_HEADER,
error_status_code,
headers_with_litellm_call_id,
with_litellm_call_id,
)
from litellm.types.token_count import TokenCountSource
from litellm.types.utils import TokenCountResponse
@ -28,6 +35,7 @@ def _logfmt_value(value: str | None) -> str:
_Counted: Final = TypeVar("_Counted")
_Rejection: Final = TypeVar("_Rejection")
_ProviderCount: Final = TypeVar("_ProviderCount", bound=Mapping[str, object])
_Answer: Final = TypeVar("_Answer")
@dataclass(frozen=True, slots=True)
@ -131,3 +139,26 @@ async def logged_provider_token_count(
)
return await _logged(count, _CountRequest(litellm_call_id, request_model, model_id), provider_fields)
async def answered_with_litellm_call_id(
answer: Awaitable[_Answer], litellm_call_id: str, fastapi_response: Response
) -> _Answer:
try:
answered: Final = await answer
except ProxyException as e:
raise with_litellm_call_id(e, litellm_call_id)
except HTTPException as e:
raise HTTPException(
status_code=e.status_code,
detail=e.detail,
headers=headers_with_litellm_call_id(e.headers, litellm_call_id),
) from e
except Exception as e:
log_llm_api_exception(e, litellm_call_id)
raise with_litellm_call_id(proxy_exception_from_route_error(e), litellm_call_id) from e
if isinstance(answered, Response):
answered.headers.setdefault(LITELLM_CALL_ID_HEADER, litellm_call_id)
return answered
fastapi_response.headers[LITELLM_CALL_ID_HEADER] = litellm_call_id # rebind-ok: FastAPI's only way to add a header
return answered

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

@ -22,7 +22,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg
read_request_body,
)
from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER
from litellm.proxy.common_utils.token_count_log import rejected_token_count
from litellm.proxy.common_utils.token_count_log import answered_with_litellm_call_id, rejected_token_count
from litellm.proxy.common_utils.token_count_settings import count_source_fields
from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
@ -327,7 +327,7 @@ async def google_stream_generate_content(
dependencies=[Depends(user_api_key_auth)],
response_model=TokenCountDetailsResponse,
)
async def google_count_tokens(request: Request, model_name: str):
async def google_count_tokens(request: Request, model_name: str, fastapi_response: Response):
"""
```json
return {
@ -342,10 +342,18 @@ async def google_count_tokens(request: Request, model_name: str):
}
```
"""
from litellm.proxy.common_utils.http_parsing_utils import read_request_body
from litellm.proxy.proxy_server import token_counter as internal_token_counter
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER))
return await answered_with_litellm_call_id(
_google_count_tokens(request, model_name, litellm_call_id), litellm_call_id, fastapi_response
)
async def _google_count_tokens(
request: Request, model_name: str, litellm_call_id: str
) -> JSONResponse | TokenCountDetailsResponse:
from litellm.proxy.common_utils.http_parsing_utils import read_request_body
from litellm.proxy.proxy_server import count_request_tokens
data: Final = await read_request_body(request=request)
try:
token_request: Final = _google_count_tokens_request(model_name, data)
@ -365,7 +373,7 @@ async def google_count_tokens(request: Request, model_name: str):
)
try:
token_response: Final = await internal_token_counter(
token_response: Final = await count_request_tokens(
request=token_request, call_endpoint=True, litellm_call_id=litellm_call_id
)
except ProxyException as error:

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

@ -107,6 +107,7 @@ from litellm.proxy.common_utils.sse_keepalive import (
wrap_passthrough_sse_bytes_with_keepalive_pings,
)
from litellm.proxy.common_utils.token_count_log import (
answered_with_litellm_call_id,
deployment_model_id,
logged_provider_token_count,
rejected_token_count,
@ -1269,7 +1270,7 @@ async def handle_bedrock_count_tokens(
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth,
request_body: dict[str, object],
) -> dict[str, object]:
) -> Mapping[str, object]:
"""
Handle AWS Bedrock CountTokens API requests.
@ -1277,11 +1278,19 @@ async def handle_bedrock_count_tokens(
- /v1/messages/count_tokens
- /v1/messages/count-tokens
"""
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER))
return await answered_with_litellm_call_id(
_bedrock_count_tokens(user_api_key_dict, request_body, litellm_call_id), litellm_call_id, fastapi_response
)
async def _bedrock_count_tokens(
user_api_key_dict: UserAPIKeyAuth, request_body: dict[str, object], litellm_call_id: str
) -> Mapping[str, object]:
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
from litellm.proxy.proxy_server import llm_router
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER))
try:
# Initialize the handler
handler: Final = BedrockCountTokensHandler()

10
litellm/proxy/proxy_server.py
View file

@ -559,6 +559,7 @@ from litellm.proxy.common_utils.timezone_utils import (
get_budget_reset_time,
)
from litellm.proxy.common_utils.token_count_log import (
answered_with_litellm_call_id,
deployment_model_id,
logged_token_count,
rejected_token_count,
@ -15306,19 +15307,26 @@ def _system_message(system: object) -> ChatCompletionSystemMessage | None:
)
async def token_counter(
request: TokenCountRequest,
fastapi_response: Response,
call_endpoint: bool = False,
litellm_call_id: Annotated[str | None, Header(alias=LITELLM_CALL_ID_HEADER)] = None,
):
"""
Args:
request: TokenCountRequest
fastapi_response: answers the x-litellm-call-id header the count's log line names
call_endpoint: bool - When set to "True" it will call the token counting endpoint - e.g Anthropic or Google AI Studio Token Counting APIs.
litellm_call_id: the x-litellm-call-id header, logged with the count so it joins the inference call's logs
Returns:
TokenCountResponse
"""
return await count_request_tokens(request=request, call_endpoint=call_endpoint, litellm_call_id=litellm_call_id)
call_id: Final = resolve_litellm_call_id(litellm_call_id)
return await answered_with_litellm_call_id(
count_request_tokens(request=request, call_endpoint=call_endpoint, litellm_call_id=call_id),
call_id,
fastapi_response,
)
async def count_request_tokens(

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

@ -70,8 +70,9 @@ from litellm.proxy.common_utils.openai_error_payload import (
LITELLM_CALL_ID_HEADER,
error_status_code,
openai_error_type_for_status,
with_litellm_call_id,
)
from litellm.proxy.common_utils.token_count_log import rejected_token_count
from litellm.proxy.common_utils.token_count_log import answered_with_litellm_call_id, rejected_token_count
from litellm.proxy.common_utils.token_count_settings import count_source_fields
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decrypted_response_id
@ -1727,6 +1728,7 @@ def _responses_input_as_token_count_messages(
)
async def responses_input_tokens(
request: Request,
fastapi_response: Response,
token_counter: _TokenCounter = _token_counter_dep,
llm_router: Router | None = _llm_router_dep,
user_api_key_dict: UserAPIKeyAuth = _user_api_key_auth_dep,
@ -1751,6 +1753,24 @@ async def responses_input_tokens(
Returns: `{"object": "response.input_tokens", "input_tokens": <count>}`
"""
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER))
return await answered_with_litellm_call_id(
_responses_input_tokens(
request, token_counter, llm_router, user_api_key_dict, response_id_security, managed_files, litellm_call_id
),
litellm_call_id,
fastapi_response,
)
async def _responses_input_tokens(
request: Request,
token_counter: _TokenCounter,
llm_router: Router | None,
user_api_key_dict: UserAPIKeyAuth,
response_id_security: ResponsesIDSecurity | None,
managed_files: _ManagedFileAccess | None,
litellm_call_id: str,
) -> JSONResponse | _ResponsesInputTokensResult:
data: Final = await read_request_body(request=request)
model_name: Final = data.get("model")
input_value: Final = data.get("input")
@ -1833,7 +1853,7 @@ async def responses_input_tokens(
except ProxyException as e:
if e.type != _PROVIDER_COUNT_ERROR_TYPE:
raise
return _provider_count_error_response(e, request)
return _provider_count_error_response(with_litellm_call_id(e, litellm_call_id), request)
result: Final[_ResponsesInputTokensResult] = {
"object": "response.input_tokens",
"input_tokens": token_response.total_tokens,

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

@ -11,6 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from fastapi import Response
from fastapi.testclient import TestClient
from litellm import Router
@ -464,7 +465,9 @@ class TestErrorLogCarriesCallId:
), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block
caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"),
):
response = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
response = await ep.count_tokens(
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
)
assert response.status_code == 500
record = self._error_record(caplog)
@ -516,7 +519,9 @@ async def test_count_tokens_answers_400_for_malformed_tools_without_counting(mod
patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts it is never reached
patch.object(proxy_server, "llm_router", new=_claude_router()), # test-quality-ok: module global imported at call time; no injection seam on this route
):
response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
response: Final = await ep.count_tokens(
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
)
assert response.status_code == 400
assert json.loads(response.body)["error"]["type"] == "invalid_request_error"
@ -539,7 +544,9 @@ async def test_count_tokens_counts_malformed_tool_entries_for_a_model_the_router
patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts what it receives
patch.object(proxy_server, "llm_router", new=_claude_router()), # test-quality-ok: module global imported at call time; no injection seam on this route
):
response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
response: Final = await ep.count_tokens(
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
)
assert response == {"input_tokens": 7}
assert counter.await_args.kwargs["request"].tools == _MALFORMED_TOOL_ENTRY
@ -572,7 +579,9 @@ async def test_count_tokens_drops_malformed_tools_when_drop_params_is_set(
patch.object(ep, "read_request_body", new=AsyncMock(return_value=body)), # test-quality-ok: endpoint reads the body via a module function; no injection seam
patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts what it receives
):
response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
response: Final = await ep.count_tokens(
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
)
assert response == {"input_tokens": 7}
counted: Final = counter.await_args.kwargs["request"]
@ -611,7 +620,9 @@ async def test_count_tokens_drops_malformed_tools_the_counted_deployment_drops_e
patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts what it receives
patch.object(proxy_server, "llm_router", new=router), # test-quality-ok: module global imported at call time; no injection seam on this route
):
response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
response: Final = await ep.count_tokens(
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
)
assert response == {"input_tokens": 7}
assert counter.await_args.kwargs["request"].tools is None
@ -781,7 +792,7 @@ def test_count_tokens_rejected_before_counting_logs_the_400_it_answers(
) -> None:
response: Final = _count_on_local_model(monkeypatch, body)
assert response.status_code == 400, response.text
assert (response.status_code, response.headers.get("x-litellm-call-id")) == (400, "count-log-local"), response.text
assert token_count_log_lines() == (
{
"litellm_call_id": "count-log-local",
@ -795,6 +806,51 @@ def test_count_tokens_rejected_before_counting_logs_the_400_it_answers(
)
def _claude_count_reply(status_code: int) -> Callable[[httpx.Request], httpx.Response]:
reply: Final = (
{"input_tokens": 77}
if status_code == 200
else {"type": "error", "error": {"type": "invalid_request_error", "message": "invalid base64 data"}}
)
return lambda request: httpx.Response(status_code, json=reply, request=request)
@pytest.mark.parametrize("client_call_id", [None, "count-echo-claude"])
@pytest.mark.parametrize("provider_status", [200, 400])
def test_count_tokens_answers_the_call_id_its_count_line_names(
monkeypatch: pytest.MonkeyPatch,
token_count_log_lines: Callable[[], tuple[dict[str, object], ...]],
provider_status: int,
client_call_id: str | None,
) -> None:
response: Final = _count_on_claude(
monkeypatch,
_claude_count_reply(provider_status),
{} if client_call_id is None else {"x-litellm-call-id": client_call_id},
)
(line,) = token_count_log_lines()
assert (response.status_code, response.headers.get("x-litellm-call-id")) == (
provider_status,
line["litellm_call_id"],
), response.text
assert client_call_id in (None, line["litellm_call_id"])
def test_count_tokens_error_body_names_the_call_id_its_header_carries_when_opted_in(
monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]]
) -> None:
import litellm.proxy.proxy_server as proxy_server
monkeypatch.setattr(proxy_server, "general_settings", {"include_call_id_in_error_body": True})
response: Final = _count_on_claude(monkeypatch, _claude_count_reply(400), {})
(line,) = token_count_log_lines()
assert response.json()["error"]["litellm_call_id"] == response.headers.get("x-litellm-call-id")
assert response.headers.get("x-litellm-call-id") == line["litellm_call_id"]
class TestEventLoggingBatchEndpoint:
"""Test the stubbed event logging batch endpoint"""

42
tests/unit/proxy/common_utils/test_token_count_log.py
View file

@ -1,9 +1,13 @@
import asyncio
from collections.abc import Callable
from typing import Final
import pytest
from fastapi import HTTPException, Response
from litellm.proxy.common_utils.token_count_log import logged_token_count
from litellm.proxy._types import ProxyException
from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER
from litellm.proxy.common_utils.token_count_log import answered_with_litellm_call_id, logged_token_count
from litellm.types.utils import TokenCountResponse
@ -43,3 +47,39 @@ async def test_a_cancelled_count_writes_no_line(
await logged_token_count(_count_ending_in(asyncio.CancelledError()), "count-log-cancelled", "count-log-gpt", None)
assert token_count_log_lines() == ()
@pytest.mark.asyncio
async def test_an_http_error_answers_its_status_and_detail_with_the_call_id_beside_its_own_headers() -> None:
rejection: Final = HTTPException(status_code=429, detail={"error": "slow down"}, headers={"retry-after": "3"})
with pytest.raises(HTTPException) as answered:
await answered_with_litellm_call_id(_count_ending_in(rejection), "count-echo-http", Response())
assert (answered.value.status_code, answered.value.detail, dict(answered.value.headers or {})) == (
429,
{"error": "slow down"},
{"retry-after": "3", LITELLM_CALL_ID_HEADER: "count-echo-http"},
)
@pytest.mark.asyncio
async def test_an_unhandled_error_answers_a_500_with_the_call_id() -> None:
with pytest.raises(ProxyException) as answered:
await answered_with_litellm_call_id(
_count_ending_in(RuntimeError("tokenizer crashed")), "count-echo-crash", Response()
)
assert (answered.value.code, answered.value.headers) == ("500", {LITELLM_CALL_ID_HEADER: "count-echo-crash"})
@pytest.mark.asyncio
async def test_a_cancelled_answer_stays_cancelled_and_names_no_call_id() -> None:
cancelled: Final = asyncio.CancelledError()
fastapi_response: Final = Response()
with pytest.raises(asyncio.CancelledError) as raised:
await answered_with_litellm_call_id(_count_ending_in(cancelled), "count-echo-cancelled", fastapi_response)
assert raised.value is cancelled
assert LITELLM_CALL_ID_HEADER not in fastapi_response.headers

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

@ -8,6 +8,7 @@ import pytest
from dotenv import load_dotenv
from pydantic import JsonValue, ValidationError
from starlette.requests import Request
from starlette.responses import Response
from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request, google_count_tokens
from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
@ -102,6 +103,7 @@ async def test_proxy_gemini_to_openai_like_model_token_counting():
}
),
model_name="volcengine/foo",
fastapi_response=Response(),
)
assert response.get("totalTokens") > 0

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

@ -172,7 +172,7 @@ def test_google_count_tokens_unchanged():
)
with patch(
"litellm.proxy.proxy_server.token_counter",
"litellm.proxy.proxy_server.count_request_tokens",
new_callable=AsyncMock,
return_value=fake_response,
):
@ -194,7 +194,7 @@ def test_google_count_tokens_unexpected_failure_uses_the_google_error_envelope()
pytest.skip(f"Skipping test due to missing dependency: {e}")
with patch(
"litellm.proxy.proxy_server.token_counter",
"litellm.proxy.proxy_server.count_request_tokens",
new_callable=AsyncMock,
side_effect=RuntimeError("tokenizer unavailable"),
):
@ -227,7 +227,7 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs()
)
with patch(
"litellm.proxy.proxy_server.token_counter",
"litellm.proxy.proxy_server.count_request_tokens",
new_callable=AsyncMock,
return_value=fake_response,
):
@ -284,6 +284,7 @@ def test_google_count_tokens_logs_one_line_with_the_callers_call_id_model_and_co
)
assert (response.status_code, response.json()["totalTokens"]) == (200, 77), response.text
assert response.headers.get("x-litellm-call-id") == "count-log-gemini"
assert token_count_log_lines() == (
{
"litellm_call_id": "count-log-gemini",
@ -318,3 +319,18 @@ def test_google_count_tokens_rejected_body_logs_the_400_it_answers(
"status_code": 400,
},
)
@pytest.mark.parametrize("client_call_id", [None, "count-echo-gemini"])
def test_google_count_tokens_rejected_body_answers_the_call_id_its_count_line_names(
token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], client_call_id: str | None
) -> None:
response: Final = _build_test_client().post(
"/v1beta/models/count-log-claude:countTokens",
json={"contents": "not a list"},
headers={} if client_call_id is None else {"x-litellm-call-id": client_call_id},
)
(line,) = token_count_log_lines()
assert (response.status_code, response.headers.get("x-litellm-call-id")) == (400, line["litellm_call_id"])
assert client_call_id in (None, line["litellm_call_id"])

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

@ -7255,7 +7255,10 @@ _BEDROCK_COUNT_MODEL: Final = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1
async def _bedrock_native_count(
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, aws_reply: httpx.Response
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
aws_reply: httpx.Response,
fastapi_response: Response,
) -> dict[str, object]:
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
handle_bedrock_count_tokens,
@ -7293,7 +7296,7 @@ async def _bedrock_native_count(
return await handle_bedrock_count_tokens(
endpoint=f"model/{_BEDROCK_COUNT_ALIAS}/count-tokens",
request=request,
fastapi_response=Response(),
fastapi_response=fastapi_response,
user_api_key_dict=MagicMock(),
request_body={"model": _BEDROCK_COUNT_ALIAS, "messages": [{"role": "user", "content": "hi"}]},
)
@ -7306,11 +7309,13 @@ async def test_bedrock_native_count_tokens_logs_the_count_against_the_router_mod
respx_mock: respx.MockRouter,
token_count_log_lines: Callable[[], tuple[dict[str, object], ...]],
) -> None:
fastapi_response: Final = Response()
result: Final = await _bedrock_native_count(
monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12})
monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12}), fastapi_response
)
assert result == {"input_tokens": 12}
assert fastapi_response.headers.get("x-litellm-call-id") == "bedrock-native-count"
assert len(respx_mock.calls) == 1
assert respx_mock.calls.last.request.url.path == "/model/anthropic.claude-haiku-4-5-20251001-v1:0/count-tokens"
assert token_count_log_lines() == (
@ -7335,9 +7340,12 @@ async def test_bedrock_native_count_tokens_logs_the_provider_status_when_the_cou
) -> None:
rejected: Final = httpx.Response(400, json={"message": "The provided model doesn't support counting tokens."})
with pytest.raises(HTTPException) as exc_info:
await _bedrock_native_count(monkeypatch, respx_mock, rejected)
await _bedrock_native_count(monkeypatch, respx_mock, rejected, Response())
assert exc_info.value.status_code == 400
assert (exc_info.value.status_code, (exc_info.value.headers or {}).get("x-litellm-call-id")) == (
400,
"bedrock-native-count",
)
assert token_count_log_lines() == (
{
"litellm_call_id": "bedrock-native-count",
@ -7391,6 +7399,33 @@ async def test_bedrock_native_count_tokens_without_a_model_logs_the_400_it_answe
@pytest.mark.asyncio
async def test_bedrock_native_count_tokens_without_a_model_answers_the_call_id_its_count_line_names(
token_count_log_lines: Callable[[], tuple[dict[str, object], ...]],
) -> None:
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
handle_bedrock_count_tokens,
)
request: Final = Request(
{"type": "http", "method": "POST", "path": f"/bedrock/model/{_BEDROCK_COUNT_ALIAS}/count-tokens", "headers": []}
)
with pytest.raises(HTTPException) as exc_info:
await handle_bedrock_count_tokens(
endpoint=f"model/{_BEDROCK_COUNT_ALIAS}/count-tokens",
request=request,
fastapi_response=Response(),
user_api_key_dict=MagicMock(),
request_body={"messages": [{"role": "user", "content": "hi"}]},
)
(line,) = token_count_log_lines()
assert (exc_info.value.status_code, (exc_info.value.headers or {}).get("x-litellm-call-id")) == (
400,
line["litellm_call_id"],
)
@pytest.mark.asyncio
async def test_bedrock_native_count_tokens_with_a_region_only_model_answers_the_validation_error_and_logs_it(
monkeypatch: pytest.MonkeyPatch,

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

@ -4830,6 +4830,7 @@ class TestResponsesInputTokens:
)
assert (response.status_code, response.json().get("input_tokens")) == (200, 7), response.text
assert response.headers.get("x-litellm-call-id") == "count-log-responses"
assert token_count_log_lines() == (
{
"litellm_call_id": "count-log-responses",
@ -4863,6 +4864,48 @@ class TestResponsesInputTokens:
},
)
def test_count_rejected_before_counting_answers_the_call_id_its_count_line_names(
self, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]]
) -> None:
response, _ = self._post_input_tokens({"model": "count-log-gpt"})
(line,) = token_count_log_lines()
assert (response.status_code, response.headers.get("x-litellm-call-id")) == (400, line["litellm_call_id"])
@pytest.mark.parametrize(
("raised", "status_code"),
[
(
ProxyException(
message="Token counting is disabled", type="token_counting_disabled", param="model", code=503
),
503,
),
(
ProxyException(
message="LiteLLM does not support counting tokens for this model.",
type="token_counting_error",
param="model",
code=400,
),
400,
),
(RuntimeError("tokenizer crashed"), 500),
],
ids=["proxy-error", "provider-count-error", "crash"],
)
def test_a_failed_count_answers_the_callers_call_id(self, raised: Exception, status_code: int) -> None:
response, _ = self._post_input_tokens(
{"model": "gpt-4o", "input": "hi"},
counter=AsyncMock(side_effect=raised),
headers={"x-litellm-call-id": "count-echo-failed"},
)
assert (response.status_code, response.headers.get("x-litellm-call-id")) == (
status_code,
"count-echo-failed",
), response.text
def test_a_managed_file_the_key_cannot_read_answers_403_without_a_count(self) -> None:
from litellm.integrations.custom_logger import CustomLogger
@ -5149,6 +5192,21 @@ class TestResponsesInputTokens:
}
def test_provider_count_error_body_names_the_call_id_its_header_carries_when_opted_in(self) -> None:
failing_counter: Final = AsyncMock(
side_effect=ProxyException(
message="LiteLLM does not support counting tokens for this model.",
type="token_counting_error",
param="model",
code="400",
)
)
with patch("litellm.proxy.proxy_server.general_settings", {"include_call_id_in_error_body": True}):
response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter)
assert response.json()["error"]["litellm_call_id"] == response.headers.get("x-litellm-call-id")
class TestGuardrailBlockedResponsesShape:
"""A pre_call block raises ModifyResponseException before any provider call.

10
tests/unit/proxy/test_custom_tokenizer_bug.py
View file

@ -19,7 +19,7 @@ import litellm.proxy.proxy_server
import litellm.utils
from litellm import Router
from litellm.proxy._types import ProxyException, TokenCountRequest
from litellm.proxy.proxy_server import token_counter
from litellm.proxy.proxy_server import count_request_tokens
def _fake_hf_tokenizer(num_tokens: int) -> MagicMock:
@ -60,7 +60,7 @@ async def test_custom_tokenizer_from_model_info_is_used(monkeypatch):
with patch.object(litellm.utils, "tokenizer_dispatch") as mock_tokenizer_cls:
mock_tokenizer_cls.from_pretrained.return_value = _fake_hf_tokenizer(7)
response = await token_counter(
response = await count_request_tokens(
request=TokenCountRequest(
model="my-embedding-model",
messages=[{"role": "user", "content": "Bonjour le monde"}],
@ -92,7 +92,7 @@ async def test_model_without_custom_tokenizer_uses_default(monkeypatch):
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
with patch.object(litellm.utils, "tokenizer_dispatch") as mock_tokenizer_cls:
response = await token_counter(
response = await count_request_tokens(
request=TokenCountRequest(
model="gpt-4",
messages=[{"role": "user", "content": "hello"}],
@ -133,7 +133,7 @@ async def test_custom_tokenizer_counts_a_vertex_partner_model_without_a_count_ro
with patch.object(litellm.utils, "tokenizer_dispatch") as mock_tokenizer_cls:
mock_tokenizer_cls.from_pretrained.return_value = _fake_hf_tokenizer(9)
response: Final = await token_counter(
response: Final = await count_request_tokens(
request=TokenCountRequest(model="vertex-glm", messages=[{"role": "user", "content": "Return exactly OK."}]),
call_endpoint=True,
)
@ -151,7 +151,7 @@ async def test_contents_only_count_of_a_vertex_partner_model_without_a_count_rou
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _vertex_glm_router_with_its_own_tokenizer())
with pytest.raises(ProxyException) as raised:
await token_counter(
await count_request_tokens(
request=TokenCountRequest(
model="vertex-glm", contents=[{"role": "user", "parts": [{"text": "Return exactly OK."}]}]
),

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

@ -36,7 +36,7 @@ load_dotenv()
# this file is to test litellm/proxy
from fastapi import Request
from fastapi import Request, Response
from fastapi.responses import JSONResponse
from starlette.types import Message
@ -78,7 +78,6 @@ from litellm.proxy.proxy_server import (
_provider_token_count,
_try_provider_token_count,
count_request_tokens,
token_counter,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted
@ -131,7 +130,7 @@ async def test_vLLM_token_counting():
setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
response = await token_counter(
response = await count_request_tokens(
request=TokenCountRequest(
model="special-alias",
messages=[{"role": "user", "content": "hello"}],
@ -164,7 +163,7 @@ async def test_token_counting_model_not_in_model_list():
setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
response = await token_counter(
response = await count_request_tokens(
request=TokenCountRequest(
model="special-alias",
messages=[{"role": "user", "content": "hello"}],
@ -197,7 +196,7 @@ async def test_gpt_token_counting():
setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
response = await token_counter(
response = await count_request_tokens(
request=TokenCountRequest(
model="gpt-4",
messages=[{"role": "user", "content": "hello"}],
@ -240,7 +239,7 @@ async def test_token_counter_leaves_the_routers_deployment_without_the_credentia
)
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
response = await token_counter(
response = await count_request_tokens(
request=TokenCountRequest(model="devstral", messages=[{"role": "user", "content": "hello"}])
)
@ -271,10 +270,10 @@ async def test_token_counter_reads_an_oidc_api_key_only_when_it_calls_the_provid
request = TokenCountRequest(model="claude", messages=[{"role": "user", "content": "hello"}])
if not call_endpoint:
assert (await token_counter(request=request, call_endpoint=call_endpoint)).total_tokens > 0
assert (await count_request_tokens(request=request, call_endpoint=call_endpoint)).total_tokens > 0
return
with pytest.raises(ProxyException) as excinfo:
await token_counter(request=request, call_endpoint=call_endpoint)
await count_request_tokens(request=request, call_endpoint=call_endpoint)
assert excinfo.value.code == "401"
assert str(tmp_path) not in excinfo.value.message
@ -389,7 +388,7 @@ async def test_anthropic_messages_count_tokens_endpoint():
try:
# Call the endpoint
response = await count_tokens(mock_request, mock_user_api_key_dict)
response = await count_tokens(mock_request, Response(), mock_user_api_key_dict)
# Verify response format matches Anthropic spec
assert isinstance(response, dict)
@ -471,7 +470,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model():
try:
# Call the endpoint
response = await count_tokens(mock_request, mock_user_api_key_dict)
response = await count_tokens(mock_request, Response(), mock_user_api_key_dict)
# Verify response format matches Anthropic spec
assert isinstance(response, dict)
@ -512,7 +511,7 @@ async def test_internal_token_counter_anthropic_provider_detection():
# Test with is_direct_request=False (simulating call from Anthropic endpoint)
# The provider counter is engaged: its API error surfaces instead of a local count
with pytest.raises(ProxyException) as exc_info:
await token_counter(
await count_request_tokens(
request=TokenCountRequest(
model="claude-test",
messages=[{"role": "user", "content": "hello"}],
@ -536,7 +535,7 @@ async def test_internal_token_counter_anthropic_provider_detection():
setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
with pytest.raises(ProxyException) as missing_key:
await token_counter(
await count_request_tokens(
request=TokenCountRequest(
model="gpt-test",
messages=[{"role": "user", "content": "hello"}],
@ -579,7 +578,7 @@ async def test_anthropic_endpoint_error_handling():
anthropic_endpoints.read_request_body = mock_read_request_body
try:
response = await count_tokens(mock_request, mock_user_api_key_dict)
response = await count_tokens(mock_request, Response(), mock_user_api_key_dict)
assert response.status_code == 400
assert json.loads(response.body) == {
@ -870,7 +869,7 @@ async def test_vertex_ai_anthropic_token_counting():
mock_handle_count_tokens.return_value = mock_token_response
# Test with messages format and call_endpoint=True
response = await token_counter(
response = await count_request_tokens(
request=TokenCountRequest(
model="vertex_ai/claude-3-5-sonnet-20241022",
messages=[
@ -1085,7 +1084,7 @@ async def test_token_counter_httpx_status_error_raises_proxy_exception():
litellm.proxy.proxy_server.llm_router = mock_router
with pytest.raises(ProxyException) as exc_info:
await token_counter(
await count_request_tokens(
request=TokenCountRequest(
model="claude-4-6-sonnet",
messages=[{"role": "user", "content": "hello"}],
@ -1151,7 +1150,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled():
litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter
with pytest.raises(ProxyException) as exc_info:
await token_counter(
await count_request_tokens(
request=TokenCountRequest(
model="claude-bedrock",
messages=[{"role": "user", "content": "hello"}],
@ -1216,7 +1215,7 @@ async def test_proxy_token_counter_error_raises_when_enabled():
with patch("litellm.token_counter") as mock_local_counter:
with pytest.raises(ProxyException) as exc_info:
await token_counter(
await count_request_tokens(
request=TokenCountRequest(
model="claude-bedrock",
messages=[{"role": "user", "content": "hello"}],
@ -1280,7 +1279,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format():
proxy_server.count_request_tokens = mock_token_counter_error
try:
response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict)
response: Final[object] = await anthropic_count_tokens(mock_request, Response(), mock_user_api_key_dict)
assert isinstance(response, JSONResponse)
assert response.status_code == 400
@ -1339,7 +1338,7 @@ async def test_anthropic_endpoint_403_permission_error_format():
proxy_server.count_request_tokens = mock_token_counter_error
try:
response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict)
response: Final[object] = await anthropic_count_tokens(mock_request, Response(), mock_user_api_key_dict)
assert isinstance(response, JSONResponse)
assert response.status_code == 403
@ -1398,7 +1397,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format():
proxy_server.count_request_tokens = mock_token_counter_error
try:
response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict)
response: Final[object] = await anthropic_count_tokens(mock_request, Response(), mock_user_api_key_dict)
assert isinstance(response, JSONResponse)
assert response.status_code == 429
@ -1474,7 +1473,7 @@ async def test_anthropic_count_tokens_drops_a_malformed_tool_result_text_when_dr
route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens").mock(
return_value=httpx.Response(200, json={"input_tokens": 9})
)
response: Final = await anthropic_count_tokens(request, UserAPIKeyAuth())
response: Final = await anthropic_count_tokens(request, Response(), UserAPIKeyAuth())
assert response["input_tokens"] == 9
sent: Final = json.loads(route.calls.last.request.content)
@ -1519,7 +1518,7 @@ async def _count_on_a_model_without_a_count_api(
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "proxy_config", proxy_server.ProxyConfig())
with respx.mock:
return await anthropic_count_tokens(request, UserAPIKeyAuth())
return await anthropic_count_tokens(request, Response(), UserAPIKeyAuth())
@pytest.mark.asyncio
@ -1600,7 +1599,7 @@ async def test_anthropic_count_tokens_on_a_claude_deployment_drops_a_malformed_t
route: Final = respx.post("https://api.anthropic.com/v1/messages/count_tokens").mock(
return_value=httpx.Response(200, json={"input_tokens": 7})
)
response: Final = await anthropic_count_tokens(request, UserAPIKeyAuth())
response: Final = await anthropic_count_tokens(request, Response(), UserAPIKeyAuth())
assert response["input_tokens"] == 7
sent: Final = json.loads(route.calls.last.request.content)
@ -2090,8 +2089,8 @@ async def test_framed_estimate_that_cannot_read_the_body_answers_the_provider_er
)
with pytest.raises(ProxyException) as raised:
await token_counter(request=TokenCountRequest(model="kimi", **unreadable_body), call_endpoint=True)
readable: Final = await token_counter(
await count_request_tokens(request=TokenCountRequest(model="kimi", **unreadable_body), call_endpoint=True)
readable: Final = await count_request_tokens(
request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]),
call_endpoint=True,
)
@ -2138,8 +2137,8 @@ async def test_framed_estimate_counts_google_contents_like_the_same_conversation
_framed_kimi(monkeypatch)
google_request: Final = _google_count_tokens_request("kimi", {"contents": contents})
from_google: Final = await token_counter(request=google_request, call_endpoint=True)
from_messages: Final = await token_counter(
from_google: Final = await count_request_tokens(request=google_request, call_endpoint=True)
from_messages: Final = await count_request_tokens(
request=TokenCountRequest(model="kimi", messages=google_request.messages), call_endpoint=True
)
@ -2170,7 +2169,7 @@ async def test_framed_deployment_whose_custom_tokenizer_cannot_load_answers_the_
)
with pytest.raises(ProxyException) as raised:
await token_counter(
await count_request_tokens(
request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]),
call_endpoint=True,
)
@ -2222,7 +2221,7 @@ def _framed_kimi_with_custom_tokenizer(monkeypatch: pytest.MonkeyPatch, identifi
async def _count_status(model: str = "kimi") -> str:
try:
response: Final = await token_counter(
response: Final = await count_request_tokens(
request=TokenCountRequest(model=model, messages=[{"role": "user", "content": "What is new?"}]),
call_endpoint=True,
)
@ -2486,8 +2485,8 @@ async def test_local_count_loads_the_tokenizer_again_after_a_failed_load(monkeyp
request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}])
with pytest.raises(OSError, match="hub unreachable"):
await token_counter(request=request)
retried: Final = await token_counter(request=request)
await count_request_tokens(request=request)
retried: Final = await count_request_tokens(request=request)
assert retried.count_source == "local_estimate" and retried.total_tokens > 0
@ -2517,15 +2516,15 @@ async def test_framed_count_loads_the_tokenizer_again_only_after_the_window_of_a
monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_MAX_RUNNING", 1)
request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}])
first: Final = await token_counter(request=request, call_endpoint=True)
second: Final = await token_counter(request=request, call_endpoint=True)
first: Final = await count_request_tokens(request=request, call_endpoint=True)
second: Final = await count_request_tokens(request=request, call_endpoint=True)
assert (first.tokenizer_type, second.tokenizer_type) == tokenizer_types
async def _fork_guard_traceback_frames() -> int:
with pytest.raises(ForkedAfterNativeRuntimeStarted) as raised:
await token_counter(
await count_request_tokens(
request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]),
call_endpoint=True,
)
@ -2564,7 +2563,7 @@ async def test_swallowed_estimate_error_is_logged_by_type_without_the_request_co
verbose_proxy_logger.addHandler(caplog.handler)
try:
with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name), pytest.raises(ProxyException):
await token_counter(
await count_request_tokens(
request=TokenCountRequest(
model="kimi",
messages=[
@ -2609,9 +2608,9 @@ async def test_local_estimate_adds_the_configured_framing_once_per_request(monke
)
messages: Final = [{"role": "system", "content": "be brief"}, {"role": "user", "content": "hello there"}]
plain: Final = await token_counter(request=TokenCountRequest(model="plain", messages=messages))
framed: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages))
framed_prompt: Final = await token_counter(request=TokenCountRequest(model="framed", prompt="hello there"))
plain: Final = await count_request_tokens(request=TokenCountRequest(model="plain", messages=messages))
framed: Final = await count_request_tokens(request=TokenCountRequest(model="framed", messages=messages))
framed_prompt: Final = await count_request_tokens(request=TokenCountRequest(model="framed", prompt="hello there"))
assert plain.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages)
assert framed.total_tokens == plain.total_tokens + 100
@ -2643,9 +2642,9 @@ async def test_framing_comes_from_the_row_the_deployment_provider_or_base_model_
messages: Final = [{"role": "user", "content": "hello there"}]
framing: Final = cost_map_entry(kimi)["token_count_framing"]
prefixed: Final = await token_counter(request=TokenCountRequest(model="prefixed", messages=messages))
split: Final = await token_counter(request=TokenCountRequest(model="split", messages=messages))
aliased: Final = await token_counter(request=TokenCountRequest(model="aliased", messages=messages))
prefixed: Final = await count_request_tokens(request=TokenCountRequest(model="prefixed", messages=messages))
split: Final = await count_request_tokens(request=TokenCountRequest(model="split", messages=messages))
aliased: Final = await count_request_tokens(request=TokenCountRequest(model="aliased", messages=messages))
assert prefixed.total_tokens == litellm.token_counter(model="system.ai.kimi-k3", messages=messages) + framing
assert split.total_tokens == prefixed.total_tokens
@ -2750,7 +2749,9 @@ async def test_bedrock_kimi_k3_deployment_with_a_custom_tokenizer_keeps_counting
monkeypatch.setattr(litellm.utils, "select_tokenizer", lambda model, custom_tokenizer=None: custom)
messages: Final = [{"role": "user", "content": "What is new?"}]
counted: Final = await token_counter(request=TokenCountRequest(model="kimi", messages=messages), call_endpoint=True)
counted: Final = await count_request_tokens(
request=TokenCountRequest(model="kimi", messages=messages), call_endpoint=True
)
assert counted.total_tokens == (
litellm.token_counter(model=counted.model_used, messages=messages, custom_tokenizer=custom)
@ -2785,7 +2786,7 @@ async def test_bedrock_kimi_k3_without_its_tokenizer_adds_the_template_offsets_t
):
_framed_kimi(monkeypatch)
counted: Final = await token_counter(
counted: Final = await count_request_tokens(
request=TokenCountRequest(model="kimi", messages=messages, tools=tools), call_endpoint=True
)
estimate: Final = litellm.token_counter(
@ -2814,7 +2815,7 @@ async def test_slow_cost_map_lookup_does_not_stall_other_requests_on_the_worker(
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None)
monkeypatch.setattr(litellm, "get_model_info", slow_get_model_info)
counting: Final = asyncio.create_task(
token_counter(
count_request_tokens(
request=TokenCountRequest(model="slow-lookup-model", messages=[{"role": "user", "content": "hi"}])
)
)
@ -2849,7 +2850,7 @@ async def test_uncountable_model_keeps_the_provider_error_when_local_counting_is
monkeypatch.setattr(litellm, "disable_token_counter", True)
with pytest.raises(ProxyException) as raised:
await token_counter(
await count_request_tokens(
request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True
)
@ -2899,7 +2900,7 @@ async def test_configured_token_count_model_overrides_the_shipped_one_and_an_emp
litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock")
)
counted: Final = await token_counter(
counted: Final = await count_request_tokens(
request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True
)
@ -2937,7 +2938,7 @@ async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geog
litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock")
)
await token_counter(
await count_request_tokens(
request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True
)
@ -2995,7 +2996,7 @@ async def test_region_lookup_for_a_shipped_count_model_leaves_the_event_loop_fre
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only"))
async def count() -> TokenCountResponse:
return await token_counter(
return await count_request_tokens(
request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]),
call_endpoint=True,
)
@ -3023,7 +3024,7 @@ async def test_framing_still_applies_while_the_deployment_is_cooling_down(monkey
)
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router)
messages: Final = [{"role": "user", "content": "hello there"}]
warm: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages))
warm: Final = await count_request_tokens(request=TokenCountRequest(model="framed", messages=messages))
router.cooldown_cache.add_deployment_to_cooldown(
model_id="framed-deployment",
@ -3031,7 +3032,7 @@ async def test_framing_still_applies_while_the_deployment_is_cooling_down(monkey
exception_status=429,
cooldown_time=60,
)
cooling: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages))
cooling: Final = await count_request_tokens(request=TokenCountRequest(model="framed", messages=messages))
assert warm.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + 100
assert cooling.total_tokens - 100 == litellm.token_counter(model="framed", messages=messages)
@ -3126,7 +3127,7 @@ async def _count_through_the_messages_endpoint(
return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)),
),
):
response: Final = await anthropic_count_tokens(mock_request, MagicMock())
response: Final = await anthropic_count_tokens(mock_request, Response(), MagicMock())
return response, peer.request_bodies
@ -3447,10 +3448,10 @@ async def test_call_endpoint_counts_a_prompt_like_the_same_text_sent_as_a_user_m
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer))
with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
from_prompt: Final = await token_counter(
from_prompt: Final = await count_request_tokens(
request=TokenCountRequest(model="gemini-counted", prompt=text), call_endpoint=True
)
from_message: Final = await token_counter(
from_message: Final = await count_request_tokens(
request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": text}]),
call_endpoint=True,
)
@ -3496,7 +3497,7 @@ async def test_call_endpoint_counts_a_prompt_on_a_vertex_ai_deployment(monkeypat
with patch.object(
VertexAIPartnerModelsTokenCounter, "handle_count_tokens_request", side_effect=_vertex_partner_word_count
):
from_prompt: Final = await token_counter(
from_prompt: Final = await count_request_tokens(
request=TokenCountRequest(model="vertex-counted", prompt=text), call_endpoint=True
)
@ -3525,7 +3526,7 @@ async def test_prompt_next_to_a_conversation_answers_400_without_drop_params(
patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client),
pytest.raises(ProxyException) as rejected,
):
await token_counter(
await count_request_tokens(
request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint
)
@ -3560,10 +3561,10 @@ async def test_drop_params_counts_a_conversation_without_the_prompt(
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer))
with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
with_prompt: Final = await token_counter(
with_prompt: Final = await count_request_tokens(
request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint
)
conversation_only: Final = await token_counter(
conversation_only: Final = await count_request_tokens(
request=TokenCountRequest(model="gemini-counted", **conversation), call_endpoint=call_endpoint
)
@ -3580,7 +3581,7 @@ async def test_call_endpoint_prompt_next_to_an_empty_conversation_counts_the_pro
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer))
with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
counted: Final = await token_counter(
counted: Final = await count_request_tokens(
request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation), call_endpoint=True
)
@ -3593,8 +3594,10 @@ async def test_local_count_of_a_prompt_next_to_empty_messages_counts_the_prompt(
) -> None:
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router())
with_empty: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX, messages=[]))
alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX))
with_empty: Final = await count_request_tokens(
request=TokenCountRequest(model="gemini-counted", prompt=_FOX, messages=[])
)
alone: Final = await count_request_tokens(request=TokenCountRequest(model="gemini-counted", prompt=_FOX))
assert with_empty.error is False
assert with_empty.total_tokens == alone.total_tokens > 0
@ -3622,13 +3625,13 @@ async def test_local_count_of_a_prompt_with_system_or_tools_counts_them_like_the
) -> None:
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router())
from_prompt: Final = await token_counter(
from_prompt: Final = await count_request_tokens(
request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation, **chat_fields)
)
from_message: Final = await token_counter(
from_message: Final = await count_request_tokens(
request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": _FOX}], **chat_fields)
)
prompt_alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX))
prompt_alone: Final = await count_request_tokens(request=TokenCountRequest(model="gemini-counted", prompt=_FOX))
assert from_prompt.error is False
assert from_prompt.total_tokens == from_message.total_tokens > prompt_alone.total_tokens
@ -3655,8 +3658,8 @@ async def test_call_endpoint_prompt_keeps_the_local_count_on_a_counter_that_need
request: Final = TokenCountRequest(model="anthropic-counted", prompt="The quick brown fox jumps over the lazy dog.")
with patch.object(anthropic_count_tokens_handler, "get_async_httpx_client", return_value=client):
with_endpoint: Final = await token_counter(request=request, call_endpoint=True)
local: Final = await token_counter(request=request, call_endpoint=False)
with_endpoint: Final = await count_request_tokens(request=request, call_endpoint=True)
local: Final = await count_request_tokens(request=request, call_endpoint=False)
assert (with_endpoint.error, with_endpoint.total_tokens, with_endpoint.tokenizer_type) == (
False,
@ -4009,8 +4012,8 @@ async def test_token_counter_counts_a_custom_azure_deployment_name_with_its_base
)
messages: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}]
named: Final = await token_counter(request=TokenCountRequest(model="named", messages=messages))
served: Final = await token_counter(request=TokenCountRequest(model="served", messages=messages))
named: Final = await count_request_tokens(request=TokenCountRequest(model="named", messages=messages))
served: Final = await count_request_tokens(request=TokenCountRequest(model="served", messages=messages))
assert named.total_tokens == served.total_tokens
@ -4124,6 +4127,33 @@ def test_utils_token_counter_failed_count_logs_its_status_in_place_of_a_count(
)
@pytest.mark.parametrize(
("body", "headers", "status_code"),
[
(_COUNTED_HELLO, {}, 200),
(_COUNTED_HELLO, {"x-litellm-call-id": "count-echo-utils"}, 200),
({"model": "count-log-gpt"}, {}, 400),
({**_COUNTED_HELLO, "prompt": "and a prompt"}, {"x-litellm-call-id": "count-echo-utils"}, 400),
],
ids=["counted", "counted-with-client-id", "no-messages", "failed-count-with-client-id"],
)
def test_utils_token_counter_answers_the_call_id_its_count_line_names(
monkeypatch: pytest.MonkeyPatch,
token_count_log_lines: Callable[[], tuple[dict[str, object], ...]],
body: Mapping[str, JsonValue],
headers: Mapping[str, str],
status_code: int,
) -> None:
response: Final = _utils_token_counter(monkeypatch, body, headers)
(line,) = token_count_log_lines()
assert (response.status_code, response.headers.get("x-litellm-call-id")) == (
status_code,
line["litellm_call_id"],
), response.text
assert headers.get("x-litellm-call-id") in (None, line["litellm_call_id"])
def test_utils_token_counter_json_line_keeps_its_own_fields_when_the_call_id_is_a_dict(
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,

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

@ -19068,6 +19068,7 @@ export interface paths {
* Token Counter
* @description Args:
* request: TokenCountRequest
* fastapi_response: answers the x-litellm-call-id header the count's log line names
* call_endpoint: bool - When set to "True" it will call the token counting endpoint - e.g Anthropic or Google AI Studio Token Counting APIs.
* litellm_call_id: the x-litellm-call-id header, logged with the count so it joins the inference call's logs
*