mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(proxy): answer x-litellm-call-id on every token count response (#51037)
* fix(proxy): answer x-litellm-call-id on every token count response The token count log line names a litellm_call_id, but count responses never carried it, so a client that sent no id of its own could not find its count's log line. Every count route (Anthropic Messages, Responses input_tokens, Gemini countTokens, Bedrock count-tokens, and /utils/token_counter) now answers the id on success and on errors, and a crash answers the JSON 500 the Anthropic and Gemini routes already used. * fix(proxy): keep the gateway's crash answers on count routes and stamp Responses opt-in bodies A crash on /utils/token_counter or /v1/responses/input_tokens goes back to the gateway's generic 500, or 503 when the database is down, now with the x-litellm-call-id header. Responses count rejections carry error.litellm_call_id when include_call_id_in_error_body is on * test(proxy): count through count_request_tokens in the token counter loop tests The token_counter route now also takes the response it sets x-litellm-call-id on, so these tests call the shared count function they exercise * fix(proxy): keep a count crash's own exception and log pre-count Responses failures The call id now reaches the crash handler through request.state instead of a wrapper exception, so logs and Sentry see the error the count raised. A Responses count that fails before counting writes its count line for any error, so its header names a line * fix(proxy): keep the token counter's injected args out of its API docs Tests the opt-in error body id on an Anthropic count the proxy rejects before counting * fix(proxy): keep a count route's own HTTP error when adding the call id The wrapper raised a new HTTPException with the same status, detail and headers, so a subclass lost its type and error tracking saw a new exception. It now adds the header to the original and raises it again --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
tavren27
GitHub
parent
a45e3b6605
commit
36c7f10c84
16 changed files with 627 additions and 157 deletions
|
|
@ -42,13 +42,14 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg
|
|||
read_request_body,
|
||||
)
|
||||
from litellm.proxy.common_utils.openai_error_payload import (
|
||||
LITELLM_CALL_ID_HEADER,
|
||||
error_status_code,
|
||||
litellm_call_id_headers,
|
||||
openai_error_param,
|
||||
openai_error_type,
|
||||
with_litellm_call_id,
|
||||
)
|
||||
from litellm.proxy.common_utils.token_count_log import rejected_token_count
|
||||
from litellm.proxy.common_utils.token_count_log import answered_with_litellm_call_id, rejected_token_count
|
||||
from litellm.proxy.common_utils.token_count_settings import count_source_fields
|
||||
from litellm.proxy.hooks.proxy_track_cost_callback import failure_row_response_cost
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
|
|
@ -336,6 +337,7 @@ def _provider_context_management(response: TokenCountResponse) -> Mapping[str, o
|
|||
)
|
||||
async def count_tokens(
|
||||
request: Request,
|
||||
fastapi_response: Response,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # Used for auth
|
||||
):
|
||||
"""
|
||||
|
|
@ -358,9 +360,17 @@ async def count_tokens(
|
|||
|
||||
Returns: {"input_tokens": <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, request, fastapi_response
|
||||
)
|
||||
|
||||
|
||||
async def _count_tokens(
|
||||
request: Request, user_api_key_dict: UserAPIKeyAuth, litellm_call_id: str
|
||||
) -> JSONResponse | Mapping[str, object]:
|
||||
from litellm.proxy.proxy_server import count_request_tokens, llm_router, proxy_config
|
||||
|
||||
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
|
||||
try:
|
||||
request_data: Final = await read_request_body(request=request)
|
||||
data: Final[dict] = {**request_data}
|
||||
|
|
@ -389,7 +399,9 @@ async def count_tokens(
|
|||
)
|
||||
if isinstance(token_request, ProxyException):
|
||||
return rejected_token_count(
|
||||
_anthropic_error_json_response(token_request, request), litellm_call_id, request_model
|
||||
_anthropic_error_json_response(with_litellm_call_id(token_request, litellm_call_id), request),
|
||||
litellm_call_id,
|
||||
request_model,
|
||||
)
|
||||
|
||||
token_response: Final = await count_request_tokens(
|
||||
|
|
@ -401,13 +413,17 @@ async def count_tokens(
|
|||
**_count_tokens_body_options(data, auto_drop_params),
|
||||
)
|
||||
except HTTPException as e:
|
||||
return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request)
|
||||
return _anthropic_error_json_response(
|
||||
with_litellm_call_id(proxy_exception_from_http_exception(e, {}), litellm_call_id), request
|
||||
)
|
||||
except ProxyException as e:
|
||||
return _anthropic_error_json_response(e, request)
|
||||
return _anthropic_error_json_response(with_litellm_call_id(e, litellm_call_id), request)
|
||||
except Exception as e:
|
||||
answered: Final = with_answered_status(e)
|
||||
log_llm_api_exception(answered, litellm_call_id)
|
||||
return _anthropic_error_json_response(proxy_exception_from_route_error(answered), request)
|
||||
return _anthropic_error_json_response(
|
||||
with_litellm_call_id(proxy_exception_from_route_error(answered), litellm_call_id), request
|
||||
)
|
||||
else:
|
||||
return {
|
||||
"input_tokens": token_response.total_tokens,
|
||||
|
|
|
|||
|
|
@ -6,10 +6,16 @@ from dataclasses import dataclass
|
|||
from types import MappingProxyType
|
||||
from typing import Final, TypeVar
|
||||
|
||||
from fastapi import status
|
||||
from fastapi import HTTPException, Request, Response, status
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy.common_utils.openai_error_payload import error_status_code
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.common_utils.openai_error_payload import (
|
||||
LITELLM_CALL_ID_HEADER,
|
||||
error_status_code,
|
||||
headers_with_litellm_call_id,
|
||||
with_litellm_call_id,
|
||||
)
|
||||
from litellm.types.token_count import TokenCountSource
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
|
|
@ -28,6 +34,7 @@ def _logfmt_value(value: str | None) -> str:
|
|||
_Counted: Final = TypeVar("_Counted")
|
||||
_Rejection: Final = TypeVar("_Rejection")
|
||||
_ProviderCount: Final = TypeVar("_ProviderCount", bound=Mapping[str, object])
|
||||
_Answer: Final = TypeVar("_Answer")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -131,3 +138,23 @@ async def logged_provider_token_count(
|
|||
)
|
||||
|
||||
return await _logged(count, _CountRequest(litellm_call_id, request_model, model_id), provider_fields)
|
||||
|
||||
|
||||
async def answered_with_litellm_call_id(
|
||||
answer: Awaitable[_Answer], litellm_call_id: str, request: Request, fastapi_response: Response
|
||||
) -> _Answer:
|
||||
try:
|
||||
answered: Final = await answer
|
||||
except ProxyException as e:
|
||||
raise with_litellm_call_id(e, litellm_call_id)
|
||||
except HTTPException as e:
|
||||
e.headers = headers_with_litellm_call_id(e.headers, litellm_call_id)
|
||||
raise
|
||||
except Exception:
|
||||
request.state.litellm_call_id = litellm_call_id # rebind-ok: the crash handler reads the id from request.state
|
||||
raise
|
||||
if isinstance(answered, Response):
|
||||
answered.headers.setdefault(LITELLM_CALL_ID_HEADER, litellm_call_id)
|
||||
return answered
|
||||
fastapi_response.headers[LITELLM_CALL_ID_HEADER] = litellm_call_id # rebind-ok: FastAPI's only way to add a header
|
||||
return answered
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg
|
|||
read_request_body,
|
||||
)
|
||||
from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER
|
||||
from litellm.proxy.common_utils.token_count_log import rejected_token_count
|
||||
from litellm.proxy.common_utils.token_count_log import answered_with_litellm_call_id, rejected_token_count
|
||||
from litellm.proxy.common_utils.token_count_settings import count_source_fields
|
||||
from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
|
||||
|
||||
|
|
@ -327,7 +327,7 @@ async def google_stream_generate_content(
|
|||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=TokenCountDetailsResponse,
|
||||
)
|
||||
async def google_count_tokens(request: Request, model_name: str):
|
||||
async def google_count_tokens(request: Request, model_name: str, fastapi_response: Response):
|
||||
"""
|
||||
```json
|
||||
return {
|
||||
|
|
@ -342,10 +342,18 @@ async def google_count_tokens(request: Request, model_name: str):
|
|||
}
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.common_utils.http_parsing_utils import read_request_body
|
||||
from litellm.proxy.proxy_server import token_counter as internal_token_counter
|
||||
|
||||
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER))
|
||||
return await answered_with_litellm_call_id(
|
||||
_google_count_tokens(request, model_name, litellm_call_id), litellm_call_id, request, fastapi_response
|
||||
)
|
||||
|
||||
|
||||
async def _google_count_tokens(
|
||||
request: Request, model_name: str, litellm_call_id: str
|
||||
) -> JSONResponse | TokenCountDetailsResponse:
|
||||
from litellm.proxy.common_utils.http_parsing_utils import read_request_body
|
||||
from litellm.proxy.proxy_server import count_request_tokens
|
||||
|
||||
data: Final = await read_request_body(request=request)
|
||||
try:
|
||||
token_request: Final = _google_count_tokens_request(model_name, data)
|
||||
|
|
@ -365,7 +373,7 @@ async def google_count_tokens(request: Request, model_name: str):
|
|||
)
|
||||
|
||||
try:
|
||||
token_response: Final = await internal_token_counter(
|
||||
token_response: Final = await count_request_tokens(
|
||||
request=token_request, call_endpoint=True, litellm_call_id=litellm_call_id
|
||||
)
|
||||
except ProxyException as error:
|
||||
|
|
|
|||
|
|
@ -133,6 +133,7 @@ from litellm.proxy.common_utils.sse_keepalive import (
|
|||
wrap_passthrough_sse_bytes_with_keepalive_pings,
|
||||
)
|
||||
from litellm.proxy.common_utils.token_count_log import (
|
||||
answered_with_litellm_call_id,
|
||||
deployment_model_id,
|
||||
logged_provider_token_count,
|
||||
rejected_token_count,
|
||||
|
|
@ -1526,7 +1527,7 @@ async def handle_bedrock_count_tokens(
|
|||
fastapi_response: Response,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request_body: dict[str, object],
|
||||
) -> dict[str, object]:
|
||||
) -> Mapping[str, object]:
|
||||
"""
|
||||
Handle AWS Bedrock CountTokens API requests.
|
||||
|
||||
|
|
@ -1534,11 +1535,22 @@ async def handle_bedrock_count_tokens(
|
|||
- /v1/messages/count_tokens
|
||||
- /v1/messages/count-tokens
|
||||
"""
|
||||
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER))
|
||||
return await answered_with_litellm_call_id(
|
||||
_bedrock_count_tokens(user_api_key_dict, request_body, litellm_call_id),
|
||||
litellm_call_id,
|
||||
request,
|
||||
fastapi_response,
|
||||
)
|
||||
|
||||
|
||||
async def _bedrock_count_tokens(
|
||||
user_api_key_dict: UserAPIKeyAuth, request_body: dict[str, object], litellm_call_id: str
|
||||
) -> Mapping[str, object]:
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER))
|
||||
try:
|
||||
# Initialize the handler
|
||||
handler: Final = BedrockCountTokensHandler()
|
||||
|
|
|
|||
|
|
@ -626,6 +626,7 @@ from litellm.proxy.common_utils.timezone_utils import (
|
|||
get_budget_reset_time,
|
||||
)
|
||||
from litellm.proxy.common_utils.token_count_log import (
|
||||
answered_with_litellm_call_id,
|
||||
deployment_model_id,
|
||||
logged_token_count,
|
||||
rejected_token_count,
|
||||
|
|
@ -2595,6 +2596,9 @@ async def otel_unhandled_exception_handler(request: Request, exc: Exception):
|
|||
raise exc
|
||||
response: Final = await _unhandled_exception_response(request, exc)
|
||||
response.headers.update(trace_response_headers(getattr(request.state, "parent_otel_span", None)))
|
||||
litellm_call_id: Final = attribute_of(request.state, "litellm_call_id")
|
||||
if isinstance(litellm_call_id, str):
|
||||
response.headers.setdefault(LITELLM_CALL_ID_HEADER, litellm_call_id)
|
||||
return response
|
||||
|
||||
|
||||
|
|
@ -16090,19 +16094,30 @@ def _system_message(system: object) -> ChatCompletionSystemMessage | None:
|
|||
)
|
||||
async def token_counter(
|
||||
request: TokenCountRequest,
|
||||
http_request: Request,
|
||||
fastapi_response: Response,
|
||||
call_endpoint: bool = False,
|
||||
litellm_call_id: Annotated[str | None, Header(alias=LITELLM_CALL_ID_HEADER)] = None,
|
||||
):
|
||||
"""
|
||||
Returns:
|
||||
TokenCountResponse
|
||||
|
||||
Args:
|
||||
request: TokenCountRequest
|
||||
call_endpoint: bool - When set to "True" it will call the token counting endpoint - e.g Anthropic or Google AI Studio Token Counting APIs.
|
||||
litellm_call_id: the x-litellm-call-id header, logged with the count so it joins the inference call's logs
|
||||
|
||||
Returns:
|
||||
TokenCountResponse
|
||||
\f
|
||||
http_request: carries the call id to the error handler when the count fails
|
||||
fastapi_response: answers the x-litellm-call-id header the count's log line names
|
||||
"""
|
||||
return await count_request_tokens(request=request, call_endpoint=call_endpoint, litellm_call_id=litellm_call_id)
|
||||
call_id: Final = resolve_litellm_call_id(litellm_call_id)
|
||||
return await answered_with_litellm_call_id(
|
||||
count_request_tokens(request=request, call_endpoint=call_endpoint, litellm_call_id=call_id),
|
||||
call_id,
|
||||
http_request,
|
||||
fastapi_response,
|
||||
)
|
||||
|
||||
|
||||
async def count_request_tokens(
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from itertools import chain
|
|||
from types import MappingProxyType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Annotated,
|
||||
Any,
|
||||
Final,
|
||||
Literal,
|
||||
|
|
@ -80,8 +81,9 @@ from litellm.proxy.common_utils.openai_error_payload import (
|
|||
LITELLM_CALL_ID_HEADER,
|
||||
error_status_code,
|
||||
openai_error_type_for_status,
|
||||
with_litellm_call_id,
|
||||
)
|
||||
from litellm.proxy.common_utils.token_count_log import rejected_token_count
|
||||
from litellm.proxy.common_utils.token_count_log import answered_with_litellm_call_id, rejected_token_count
|
||||
from litellm.proxy.common_utils.token_count_settings import count_source_fields
|
||||
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
|
||||
from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decrypted_response_id
|
||||
|
|
@ -1575,6 +1577,12 @@ def _proxy_managed_files_hook() -> _ManagedFileAccess | None:
|
|||
_managed_files_hook_dep: Final = Depends(_proxy_managed_files_hook)
|
||||
|
||||
|
||||
def _proxy_general_settings() -> Mapping[str, object]:
|
||||
from litellm.proxy.proxy_server import general_settings_view
|
||||
|
||||
return general_settings_view()
|
||||
|
||||
|
||||
async def _with_managed_file_id_mapping(
|
||||
data: Mapping[str, object], managed_files: _ManagedFileAccess | None, user_api_key_dict: UserAPIKeyAuth
|
||||
) -> Mapping[str, object]:
|
||||
|
|
@ -1607,7 +1615,9 @@ def _with_authorized_previous_response(
|
|||
return {**data, "previous_response_id": authorized_id, ADDRESSED_RESPONSE_ID_FIELD: previous_response_id}
|
||||
|
||||
|
||||
def _responses_invalid_request_response(message: str, param: str | None, code: str | None) -> JSONResponse:
|
||||
def _responses_invalid_request_response(
|
||||
message: str, param: str | None, code: str | None, body_call_id: str | None
|
||||
) -> JSONResponse:
|
||||
body: Final[_ResponsesApiErrorBody] = {
|
||||
"error": {
|
||||
"message": message,
|
||||
|
|
@ -1616,7 +1626,7 @@ def _responses_invalid_request_response(message: str, param: str | None, code: s
|
|||
"code": code,
|
||||
}
|
||||
}
|
||||
return JSONResponse(status_code=400, content=body)
|
||||
return JSONResponse(status_code=400, content={"error": with_call_id({**body["error"]}, body_call_id)})
|
||||
|
||||
|
||||
class _ProviderErrorMessage(TypedDict):
|
||||
|
|
@ -1681,16 +1691,17 @@ def _provider_count_error_body(raw_message: str, status_code: int) -> _Responses
|
|||
}
|
||||
|
||||
|
||||
def _provider_count_error_response(exc: ProxyException, request: Request) -> JSONResponse:
|
||||
def _provider_count_error_response(
|
||||
exc: ProxyException, request: Request, general_settings: Mapping[str, object]
|
||||
) -> JSONResponse:
|
||||
from litellm.proxy.proxy_server import (
|
||||
_close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; a route that answers an error itself must stamp the OTel server span like the global ProxyException handler does
|
||||
general_settings_view,
|
||||
)
|
||||
|
||||
status_code: Final = error_status_code(exc, 500)
|
||||
_close_dangling_otel_server_span(request, status_code, exc=exc)
|
||||
body: Final = _provider_count_error_body(exc.message, status_code)
|
||||
call_id: Final = error_body_call_id(general_settings_view(), exc.headers.get(LITELLM_CALL_ID_HEADER))
|
||||
call_id: Final = error_body_call_id(general_settings, exc.headers.get(LITELLM_CALL_ID_HEADER))
|
||||
return JSONResponse(
|
||||
status_code=status_code,
|
||||
content={"error": with_call_id({**body["error"]}, call_id)},
|
||||
|
|
@ -1698,11 +1709,12 @@ def _provider_count_error_response(exc: ProxyException, request: Request) -> JSO
|
|||
)
|
||||
|
||||
|
||||
def _missing_responses_param_response(param: str) -> JSONResponse:
|
||||
def _missing_responses_param_response(param: str, body_call_id: str | None) -> JSONResponse:
|
||||
return _responses_invalid_request_response(
|
||||
message=f"Missing required parameter: '{param}'.",
|
||||
param=param,
|
||||
code="missing_required_parameter",
|
||||
body_call_id=body_call_id,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1744,6 +1756,8 @@ def _responses_input_as_token_count_messages(
|
|||
)
|
||||
async def responses_input_tokens(
|
||||
request: Request,
|
||||
fastapi_response: Response,
|
||||
general_settings: Annotated[Mapping[str, object], Depends(_proxy_general_settings)],
|
||||
token_counter: _TokenCounter = _token_counter_dep,
|
||||
llm_router: Router | None = _llm_router_dep,
|
||||
user_api_key_dict: UserAPIKeyAuth = _user_api_key_auth_dep,
|
||||
|
|
@ -1768,19 +1782,50 @@ async def responses_input_tokens(
|
|||
Returns: `{"object": "response.input_tokens", "input_tokens": <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,
|
||||
general_settings,
|
||||
),
|
||||
litellm_call_id,
|
||||
request,
|
||||
fastapi_response,
|
||||
)
|
||||
|
||||
|
||||
async def _responses_input_tokens(
|
||||
request: Request,
|
||||
token_counter: _TokenCounter,
|
||||
llm_router: Router | None,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response_id_security: ResponsesIDSecurity | None,
|
||||
managed_files: _ManagedFileAccess | None,
|
||||
litellm_call_id: str,
|
||||
general_settings: Mapping[str, object],
|
||||
) -> JSONResponse | _ResponsesInputTokensResult:
|
||||
body_call_id: Final = error_body_call_id(general_settings, litellm_call_id)
|
||||
data: Final = await read_request_body(request=request)
|
||||
model_name: Final = data.get("model")
|
||||
input_value: Final = data.get("input")
|
||||
if not isinstance(model_name, str) or not model_name:
|
||||
return rejected_token_count(_missing_responses_param_response("model"), litellm_call_id, None)
|
||||
return rejected_token_count(_missing_responses_param_response("model", body_call_id), litellm_call_id, None)
|
||||
if input_value is None:
|
||||
return rejected_token_count(_missing_responses_param_response("input"), litellm_call_id, model_name)
|
||||
return rejected_token_count(
|
||||
_missing_responses_param_response("input", body_call_id), litellm_call_id, model_name
|
||||
)
|
||||
if isinstance(input_value, (str, list)) and not input_value:
|
||||
return rejected_token_count(
|
||||
_responses_invalid_request_response(
|
||||
message="""One of "input" or "previous_response_id" or 'prompt' or 'conversation' must be provided.""",
|
||||
param=None,
|
||||
code="missing_required_parameter",
|
||||
body_call_id=body_call_id,
|
||||
),
|
||||
litellm_call_id,
|
||||
model_name,
|
||||
|
|
@ -1792,7 +1837,9 @@ async def responses_input_tokens(
|
|||
token_fields: Final = resolve_counted_tool_fields(counted_fields, llm_router, user_api_key_dict.team_id)
|
||||
if isinstance(token_fields, ProxyException):
|
||||
return rejected_token_count(
|
||||
_responses_invalid_request_response(message=token_fields.message, param=token_fields.param, code=None),
|
||||
_responses_invalid_request_response(
|
||||
message=token_fields.message, param=token_fields.param, code=None, body_call_id=body_call_id
|
||||
),
|
||||
litellm_call_id,
|
||||
model_name,
|
||||
)
|
||||
|
|
@ -1807,7 +1854,10 @@ async def responses_input_tokens(
|
|||
except ValidationError as e:
|
||||
return rejected_token_count(
|
||||
_responses_invalid_request_response(
|
||||
message=first_validation_error(e), param=str(e.errors(include_url=False)[0]["loc"][0]), code=None
|
||||
message=first_validation_error(e),
|
||||
param=str(e.errors(include_url=False)[0]["loc"][0]),
|
||||
code=None,
|
||||
body_call_id=body_call_id,
|
||||
),
|
||||
litellm_call_id,
|
||||
model_name,
|
||||
|
|
@ -1815,7 +1865,7 @@ async def responses_input_tokens(
|
|||
except Exception as e:
|
||||
return rejected_token_count(
|
||||
_responses_invalid_request_response(
|
||||
message=f"Invalid request for token counting: {e}", param=None, code=None
|
||||
message=f"Invalid request for token counting: {e}", param=None, code=None, body_call_id=body_call_id
|
||||
),
|
||||
litellm_call_id,
|
||||
model_name,
|
||||
|
|
@ -1835,8 +1885,9 @@ async def responses_input_tokens(
|
|||
managed_files,
|
||||
user_api_key_dict,
|
||||
)
|
||||
except ProxyException as e:
|
||||
raise rejected_token_count(e, litellm_call_id, model_name)
|
||||
except Exception as e:
|
||||
rejected_token_count(e, litellm_call_id, model_name)
|
||||
raise
|
||||
try:
|
||||
token_response: Final = await token_counter(
|
||||
request=token_request,
|
||||
|
|
@ -1847,7 +1898,7 @@ async def responses_input_tokens(
|
|||
except ProxyException as e:
|
||||
if e.type != _PROVIDER_COUNT_ERROR_TYPE:
|
||||
raise
|
||||
return _provider_count_error_response(e, request)
|
||||
return _provider_count_error_response(with_litellm_call_id(e, litellm_call_id), request, general_settings)
|
||||
result: Final[_ResponsesInputTokensResult] = {
|
||||
"object": "response.input_tokens",
|
||||
"input_tokens": token_response.total_tokens,
|
||||
|
|
|
|||
|
|
@ -557,7 +557,9 @@ class TestProxyExceptionAnthropicEnvelope:
|
|||
proxy_server, "count_request_tokens", new=AsyncMock(side_effect=unavailable)
|
||||
),
|
||||
):
|
||||
response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
|
||||
response: Final = await ep.count_tokens(
|
||||
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
|
||||
)
|
||||
|
||||
body: Final = json.loads(response.body)
|
||||
assert response.status_code == expected_status
|
||||
|
|
@ -840,7 +842,9 @@ class TestErrorLogCarriesCallId:
|
|||
), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block
|
||||
caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"),
|
||||
):
|
||||
response = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
|
||||
response = await ep.count_tokens(
|
||||
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
|
||||
)
|
||||
|
||||
assert response.status_code == 500
|
||||
record = self._error_record(caplog)
|
||||
|
|
@ -883,7 +887,9 @@ class TestCountTokensCompactedHistory:
|
|||
patch.object(ep, "read_request_body", new=AsyncMock(return_value=body)), # test-quality-ok: endpoint reads the body via a module function; no injection seam
|
||||
patch.object(proxy_server, "llm_router", None), # test-quality-ok: module global; with no router the count runs on the local tokenizer, the path that raised
|
||||
):
|
||||
return await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
|
||||
return await ep.count_tokens(
|
||||
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
|
||||
)
|
||||
|
||||
counted = await count(compacted)
|
||||
|
||||
|
|
@ -934,7 +940,9 @@ async def test_count_tokens_answers_400_for_malformed_tools_without_counting(mod
|
|||
patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts it is never reached
|
||||
patch.object(proxy_server, "llm_router", new=_claude_router()), # test-quality-ok: module global imported at call time; no injection seam on this route
|
||||
):
|
||||
response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
|
||||
response: Final = await ep.count_tokens(
|
||||
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert json.loads(response.body)["error"]["type"] == "invalid_request_error"
|
||||
|
|
@ -957,7 +965,9 @@ async def test_count_tokens_counts_malformed_tool_entries_for_a_model_the_router
|
|||
patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts what it receives
|
||||
patch.object(proxy_server, "llm_router", new=_claude_router()), # test-quality-ok: module global imported at call time; no injection seam on this route
|
||||
):
|
||||
response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
|
||||
response: Final = await ep.count_tokens(
|
||||
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
|
||||
)
|
||||
|
||||
assert response == {"input_tokens": 7}
|
||||
assert counter.await_args.kwargs["request"].tools == _MALFORMED_TOOL_ENTRY
|
||||
|
|
@ -990,7 +1000,9 @@ async def test_count_tokens_drops_malformed_tools_when_drop_params_is_set(
|
|||
patch.object(ep, "read_request_body", new=AsyncMock(return_value=body)), # test-quality-ok: endpoint reads the body via a module function; no injection seam
|
||||
patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts what it receives
|
||||
):
|
||||
response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
|
||||
response: Final = await ep.count_tokens(
|
||||
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
|
||||
)
|
||||
|
||||
assert response == {"input_tokens": 7}
|
||||
counted: Final = counter.await_args.kwargs["request"]
|
||||
|
|
@ -1017,7 +1029,9 @@ async def test_count_tokens_keeps_claude_code_malformed_tools_when_operator_sets
|
|||
patch.object(proxy_server, "llm_router", new=_claude_router()), # test-quality-ok: module global imported at call time; no injection seam on this route
|
||||
patch.object(proxy_server, "proxy_config", new=operator_config), # test-quality-ok: module global imported at call time; no injection seam on this route
|
||||
):
|
||||
response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
|
||||
response: Final = await ep.count_tokens(
|
||||
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert json.loads(response.body)["error"]["type"] == "invalid_request_error"
|
||||
|
|
@ -1055,7 +1069,9 @@ async def test_count_tokens_drops_malformed_tools_the_counted_deployment_drops_e
|
|||
patch.object(proxy_server, "count_request_tokens", new=counter), # test-quality-ok: module global imported at call time; the test asserts what it receives
|
||||
patch.object(proxy_server, "llm_router", new=router), # test-quality-ok: module global imported at call time; no injection seam on this route
|
||||
):
|
||||
response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
|
||||
response: Final = await ep.count_tokens(
|
||||
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
|
||||
)
|
||||
|
||||
assert response == {"input_tokens": 7}
|
||||
assert counter.await_args.kwargs["request"].tools is None
|
||||
|
|
@ -1246,7 +1262,7 @@ def test_count_tokens_rejected_before_counting_logs_the_400_it_answers(
|
|||
) -> None:
|
||||
response: Final = _count_on_local_model(monkeypatch, body)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert (response.status_code, response.headers.get("x-litellm-call-id")) == (400, "count-log-local"), response.text
|
||||
assert token_count_log_lines() == (
|
||||
{
|
||||
"litellm_call_id": "count-log-local",
|
||||
|
|
@ -1303,6 +1319,59 @@ def test_count_tokens_on_bedrock_answers_400_for_a_message_without_content_witho
|
|||
assert len(respx_mock.calls) == 0
|
||||
|
||||
|
||||
def _claude_count_reply(status_code: int) -> Callable[[httpx.Request], httpx.Response]:
|
||||
reply: Final = (
|
||||
{"input_tokens": 77}
|
||||
if status_code == 200
|
||||
else {"type": "error", "error": {"type": "invalid_request_error", "message": "invalid base64 data"}}
|
||||
)
|
||||
return lambda request: httpx.Response(status_code, json=reply, request=request)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client_call_id", [None, "count-echo-claude"])
|
||||
@pytest.mark.parametrize("provider_status", [200, 400])
|
||||
def test_count_tokens_answers_the_call_id_its_count_line_names(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
token_count_log_lines: Callable[[], tuple[dict[str, object], ...]],
|
||||
provider_status: int,
|
||||
client_call_id: str | None,
|
||||
) -> None:
|
||||
response: Final = _count_on_claude(
|
||||
monkeypatch,
|
||||
_claude_count_reply(provider_status),
|
||||
{} if client_call_id is None else {"x-litellm-call-id": client_call_id},
|
||||
)
|
||||
|
||||
(line,) = token_count_log_lines()
|
||||
assert (response.status_code, response.headers.get("x-litellm-call-id")) == (
|
||||
provider_status,
|
||||
line["litellm_call_id"],
|
||||
), response.text
|
||||
assert client_call_id in (None, line["litellm_call_id"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("provider_status", "extra_body"),
|
||||
[(400, MappingProxyType({})), (200, MappingProxyType({"tools": 5}))],
|
||||
ids=["provider_rejection", "malformed_tools_rejection"],
|
||||
)
|
||||
def test_count_tokens_error_body_names_the_call_id_its_header_carries_when_opted_in(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
token_count_log_lines: Callable[[], tuple[dict[str, object], ...]],
|
||||
provider_status: int,
|
||||
extra_body: Mapping[str, object],
|
||||
) -> None:
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"include_call_id_in_error_body": True})
|
||||
|
||||
response: Final = _count_on_claude(monkeypatch, _claude_count_reply(provider_status), {}, extra_body)
|
||||
|
||||
(line,) = token_count_log_lines()
|
||||
assert response.json()["error"]["litellm_call_id"] == response.headers.get("x-litellm-call-id")
|
||||
assert response.headers.get("x-litellm-call-id") == line["litellm_call_id"]
|
||||
|
||||
|
||||
class TestEventLoggingBatchEndpoint:
|
||||
"""Test the stubbed event logging batch endpoint"""
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,12 @@
|
|||
import asyncio
|
||||
from collections.abc import Callable
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException, Request, Response
|
||||
|
||||
from litellm.proxy.common_utils.token_count_log import logged_token_count
|
||||
from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER
|
||||
from litellm.proxy.common_utils.token_count_log import answered_with_litellm_call_id, logged_token_count
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
|
||||
|
|
@ -15,6 +18,10 @@ async def _count_ending_in(error: BaseException) -> TokenCountResponse:
|
|||
raise error
|
||||
|
||||
|
||||
def _count_request() -> Request:
|
||||
return Request({"type": "http", "method": "POST", "path": "/utils/token_counter", "headers": []})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_count_ending_in_a_base_exception_logs_a_500(
|
||||
token_count_log_lines: Callable[[], tuple[dict[str, object], ...]],
|
||||
|
|
@ -43,3 +50,46 @@ async def test_a_cancelled_count_writes_no_line(
|
|||
await logged_token_count(_count_ending_in(asyncio.CancelledError()), "count-log-cancelled", "count-log-gpt", None)
|
||||
|
||||
assert token_count_log_lines() == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_http_error_answers_its_status_and_detail_with_the_call_id_beside_its_own_headers() -> None:
|
||||
rejection: Final = HTTPException(status_code=429, detail={"error": "slow down"}, headers={"retry-after": "3"})
|
||||
|
||||
with pytest.raises(HTTPException) as answered:
|
||||
await answered_with_litellm_call_id(_count_ending_in(rejection), "count-echo-http", _count_request(), Response())
|
||||
|
||||
assert (answered.value.status_code, answered.value.detail, dict(answered.value.headers or {})) == (
|
||||
429,
|
||||
{"error": "slow down"},
|
||||
{"retry-after": "3", LITELLM_CALL_ID_HEADER: "count-echo-http"},
|
||||
)
|
||||
|
||||
|
||||
class _CountTooLarge(HTTPException):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_http_error_reaches_its_handler_as_the_same_error_with_the_call_id() -> None:
|
||||
too_large: Final = _CountTooLarge(status_code=413, detail="count body too large")
|
||||
|
||||
with pytest.raises(_CountTooLarge) as answered:
|
||||
await answered_with_litellm_call_id(_count_ending_in(too_large), "count-echo-same", _count_request(), Response())
|
||||
|
||||
assert answered.value is too_large
|
||||
assert dict(answered.value.headers or {}) == {LITELLM_CALL_ID_HEADER: "count-echo-same"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_cancelled_answer_stays_cancelled_and_names_no_call_id() -> None:
|
||||
cancelled: Final = asyncio.CancelledError()
|
||||
fastapi_response: Final = Response()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError) as raised:
|
||||
await answered_with_litellm_call_id(
|
||||
_count_ending_in(cancelled), "count-echo-cancelled", _count_request(), fastapi_response
|
||||
)
|
||||
|
||||
assert raised.value is cancelled
|
||||
assert LITELLM_CALL_ID_HEADER not in fastapi_response.headers
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"])
|
||||
|
|
|
|||
|
|
@ -8250,6 +8250,7 @@ async def _bedrock_native_count(
|
|||
monkeypatch: pytest.MonkeyPatch,
|
||||
respx_mock: respx.MockRouter,
|
||||
aws_reply: httpx.Response,
|
||||
fastapi_response: Response,
|
||||
request_body: Mapping[str, object] = _BEDROCK_COUNT_BODY,
|
||||
) -> dict[str, object]:
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
|
|
@ -8288,7 +8289,7 @@ async def _bedrock_native_count(
|
|||
return await handle_bedrock_count_tokens(
|
||||
endpoint=f"model/{_BEDROCK_COUNT_ALIAS}/count-tokens",
|
||||
request=request,
|
||||
fastapi_response=Response(),
|
||||
fastapi_response=fastapi_response,
|
||||
user_api_key_dict=MagicMock(),
|
||||
request_body=dict(request_body),
|
||||
)
|
||||
|
|
@ -8301,9 +8302,13 @@ async def test_bedrock_native_count_tokens_logs_the_count_against_the_router_mod
|
|||
respx_mock: respx.MockRouter,
|
||||
token_count_log_lines: Callable[[], tuple[dict[str, object], ...]],
|
||||
) -> None:
|
||||
result: Final = await _bedrock_native_count(monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12}))
|
||||
fastapi_response: Final = Response()
|
||||
result: Final = await _bedrock_native_count(
|
||||
monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12}), fastapi_response
|
||||
)
|
||||
|
||||
assert result == {"input_tokens": 12}
|
||||
assert fastapi_response.headers.get("x-litellm-call-id") == "bedrock-native-count"
|
||||
assert len(respx_mock.calls) == 1
|
||||
assert respx_mock.calls.last.request.url.path == "/model/anthropic.claude-haiku-4-5-20251001-v1:0/count-tokens"
|
||||
assert token_count_log_lines() == (
|
||||
|
|
@ -8328,9 +8333,12 @@ async def test_bedrock_native_count_tokens_logs_the_provider_status_when_the_cou
|
|||
) -> None:
|
||||
rejected: Final = httpx.Response(400, json={"message": "The provided model doesn't support counting tokens."})
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _bedrock_native_count(monkeypatch, respx_mock, rejected)
|
||||
await _bedrock_native_count(monkeypatch, respx_mock, rejected, Response())
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert (exc_info.value.status_code, (exc_info.value.headers or {}).get("x-litellm-call-id")) == (
|
||||
400,
|
||||
"bedrock-native-count",
|
||||
)
|
||||
assert token_count_log_lines() == (
|
||||
{
|
||||
"litellm_call_id": "bedrock-native-count",
|
||||
|
|
@ -8363,7 +8371,7 @@ async def test_bedrock_native_count_tokens_answers_400_for_a_malformed_body_with
|
|||
) -> None:
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _bedrock_native_count(
|
||||
monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12}), request_body
|
||||
monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12}), Response(), request_body
|
||||
)
|
||||
|
||||
assert (exc_info.value.status_code, exc_info.value.detail) == (400, {"error": problem})
|
||||
|
|
@ -8458,6 +8466,33 @@ async def test_bedrock_native_count_tokens_without_a_model_logs_the_400_it_answe
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_native_count_tokens_without_a_model_answers_the_call_id_its_count_line_names(
|
||||
token_count_log_lines: Callable[[], tuple[dict[str, object], ...]],
|
||||
) -> None:
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
handle_bedrock_count_tokens,
|
||||
)
|
||||
|
||||
request: Final = Request(
|
||||
{"type": "http", "method": "POST", "path": f"/bedrock/model/{_BEDROCK_COUNT_ALIAS}/count-tokens", "headers": []}
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await handle_bedrock_count_tokens(
|
||||
endpoint=f"model/{_BEDROCK_COUNT_ALIAS}/count-tokens",
|
||||
request=request,
|
||||
fastapi_response=Response(),
|
||||
user_api_key_dict=MagicMock(),
|
||||
request_body={"messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
|
||||
(line,) = token_count_log_lines()
|
||||
assert (exc_info.value.status_code, (exc_info.value.headers or {}).get("x-litellm-call-id")) == (
|
||||
400,
|
||||
line["litellm_call_id"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_native_count_tokens_with_a_region_only_model_answers_the_validation_error_and_logs_it(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
|
|
|||
|
|
@ -5416,10 +5416,13 @@ class TestResponsesInputTokens:
|
|||
response_id_security: object | None = None,
|
||||
managed_files: object | None = None,
|
||||
headers: Mapping[str, str] = EMPTY_MAPPING,
|
||||
general_settings: Mapping[str, object] | None = None,
|
||||
raise_server_exceptions: bool = False,
|
||||
) -> tuple[Response, AsyncMock]:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.response_api_endpoints.endpoints import (
|
||||
_proxy_general_settings,
|
||||
_proxy_llm_router,
|
||||
_proxy_managed_files_hook,
|
||||
_proxy_response_id_security,
|
||||
|
|
@ -5446,8 +5449,10 @@ class TestResponsesInputTokens:
|
|||
app.dependency_overrides[_proxy_response_id_security] = lambda: response_id_security
|
||||
if managed_files is not None:
|
||||
app.dependency_overrides[_proxy_managed_files_hook] = lambda: managed_files
|
||||
if general_settings is not None:
|
||||
app.dependency_overrides[_proxy_general_settings] = lambda: general_settings
|
||||
try:
|
||||
client = TestClient(app)
|
||||
client = TestClient(app, raise_server_exceptions=raise_server_exceptions)
|
||||
response = client.post(path, json=body, headers={"Authorization": "Bearer sk-9876", **headers})
|
||||
return response, token_counter_mock
|
||||
finally:
|
||||
|
|
@ -5456,6 +5461,7 @@ class TestResponsesInputTokens:
|
|||
app.dependency_overrides.pop(_proxy_llm_router, None)
|
||||
app.dependency_overrides.pop(_proxy_response_id_security, None)
|
||||
app.dependency_overrides.pop(_proxy_managed_files_hook, None)
|
||||
app.dependency_overrides.pop(_proxy_general_settings, None)
|
||||
|
||||
def test_string_input_returns_openai_input_tokens_shape(self):
|
||||
response, counter = self._post_input_tokens({"model": "gpt-4o", "input": "Hello, how are you?"})
|
||||
|
|
@ -5738,6 +5744,7 @@ class TestResponsesInputTokens:
|
|||
)
|
||||
|
||||
assert (response.status_code, response.json().get("input_tokens")) == (200, 7), response.text
|
||||
assert response.headers.get("x-litellm-call-id") == "count-log-responses"
|
||||
assert token_count_log_lines() == (
|
||||
{
|
||||
"litellm_call_id": "count-log-responses",
|
||||
|
|
@ -5771,6 +5778,81 @@ class TestResponsesInputTokens:
|
|||
},
|
||||
)
|
||||
|
||||
def test_count_rejected_before_counting_answers_the_call_id_its_count_line_names(
|
||||
self, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]]
|
||||
) -> None:
|
||||
response, _ = self._post_input_tokens({"model": "count-log-gpt"})
|
||||
|
||||
(line,) = token_count_log_lines()
|
||||
assert (response.status_code, response.headers.get("x-litellm-call-id")) == (400, line["litellm_call_id"])
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("raised", "status_code"),
|
||||
[
|
||||
(
|
||||
ProxyException(
|
||||
message="Token counting is disabled", type="token_counting_disabled", param="model", code=503
|
||||
),
|
||||
503,
|
||||
),
|
||||
(
|
||||
ProxyException(
|
||||
message="LiteLLM does not support counting tokens for this model.",
|
||||
type="token_counting_error",
|
||||
param="model",
|
||||
code=400,
|
||||
),
|
||||
400,
|
||||
),
|
||||
],
|
||||
ids=["proxy-error", "provider-count-error"],
|
||||
)
|
||||
def test_a_failed_count_answers_the_callers_call_id(self, raised: Exception, status_code: int) -> None:
|
||||
response, _ = self._post_input_tokens(
|
||||
{"model": "gpt-4o", "input": "hi"},
|
||||
counter=AsyncMock(side_effect=raised),
|
||||
headers={"x-litellm-call-id": "count-echo-failed"},
|
||||
)
|
||||
|
||||
assert (response.status_code, response.headers.get("x-litellm-call-id")) == (
|
||||
status_code,
|
||||
"count-echo-failed",
|
||||
), response.text
|
||||
|
||||
def test_a_crashed_count_answers_the_gateways_internal_error_with_the_callers_call_id(self) -> None:
|
||||
response, _ = self._post_input_tokens(
|
||||
{"model": "gpt-4o", "input": "hi"},
|
||||
counter=AsyncMock(side_effect=RuntimeError("tokenizer crashed reading /etc/litellm/config.yaml")),
|
||||
headers={"x-litellm-call-id": "count-echo-crashed"},
|
||||
)
|
||||
|
||||
assert (response.status_code, response.headers.get("x-litellm-call-id")) == (500, "count-echo-crashed")
|
||||
assert response.json() == {"error": {"message": "Internal server error", "type": "internal_server_error"}}
|
||||
|
||||
def test_a_crashed_count_reaches_error_tracking_as_the_error_it_raised(self) -> None:
|
||||
crash: Final = RuntimeError("tokenizer crashed")
|
||||
|
||||
with pytest.raises(RuntimeError) as raised:
|
||||
self._post_input_tokens(
|
||||
{"model": "gpt-4o", "input": "hi"},
|
||||
counter=AsyncMock(side_effect=crash),
|
||||
raise_server_exceptions=True,
|
||||
)
|
||||
|
||||
assert raised.value is crash
|
||||
|
||||
def test_a_count_during_a_database_outage_answers_503_with_the_callers_call_id(self) -> None:
|
||||
from prisma.errors import ClientNotConnectedError
|
||||
|
||||
response, _ = self._post_input_tokens(
|
||||
{"model": "gpt-4o", "input": "hi"},
|
||||
counter=AsyncMock(side_effect=ClientNotConnectedError()),
|
||||
headers={"x-litellm-call-id": "count-echo-db-down"},
|
||||
)
|
||||
|
||||
assert (response.status_code, response.headers.get("x-litellm-call-id")) == (503, "count-echo-db-down")
|
||||
assert response.json()["error"]["type"] == "no_db_connection"
|
||||
|
||||
def test_a_managed_file_the_key_cannot_read_answers_403_without_a_count(self) -> None:
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
|
|
@ -5792,6 +5874,28 @@ class TestResponsesInputTokens:
|
|||
assert response.json()["error"]["message"] == denied
|
||||
counter.assert_not_awaited()
|
||||
|
||||
def test_a_managed_file_lookup_crash_answers_the_call_id_its_count_line_names(
|
||||
self, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]]
|
||||
) -> None:
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
class ManagedFiles(CustomLogger):
|
||||
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
||||
raise RuntimeError("managed file table unreadable")
|
||||
|
||||
response, counter = self._post_input_tokens(
|
||||
{
|
||||
"model": "count-log-gpt",
|
||||
"input": [{"role": "user", "content": [{"type": "input_file", "file_id": "litellm-managed-file"}]}],
|
||||
},
|
||||
managed_files=ManagedFiles(),
|
||||
)
|
||||
|
||||
(line,) = token_count_log_lines()
|
||||
assert (response.status_code, response.headers.get("x-litellm-call-id")) == (500, line["litellm_call_id"])
|
||||
assert line["status_code"] == 500
|
||||
counter.assert_not_awaited()
|
||||
|
||||
def test_malformed_tools_follow_the_counted_deployment_even_when_fallbacks_are_configured(self) -> None:
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
|
|
@ -6057,6 +6161,34 @@ class TestResponsesInputTokens:
|
|||
}
|
||||
|
||||
|
||||
def test_provider_count_error_body_names_the_call_id_its_header_carries_when_opted_in(self) -> None:
|
||||
failing_counter: Final = AsyncMock(
|
||||
side_effect=ProxyException(
|
||||
message="LiteLLM does not support counting tokens for this model.",
|
||||
type="token_counting_error",
|
||||
param="model",
|
||||
code="400",
|
||||
)
|
||||
)
|
||||
with patch("litellm.proxy.proxy_server.general_settings", {"include_call_id_in_error_body": True}):
|
||||
response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter)
|
||||
|
||||
assert response.json()["error"]["litellm_call_id"] == response.headers.get("x-litellm-call-id")
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
[{"input": "hi"}, {"model": "gpt-4o"}, {"model": "gpt-4o", "input": ""}],
|
||||
)
|
||||
def test_count_rejection_body_names_the_call_id_its_header_carries_when_opted_in(
|
||||
self, body: Mapping[str, object]
|
||||
) -> None:
|
||||
response, counter = self._post_input_tokens(dict(body), general_settings={"include_call_id_in_error_body": True})
|
||||
|
||||
counter.assert_not_awaited()
|
||||
assert response.status_code == 400, response.text
|
||||
assert response.json()["error"]["litellm_call_id"] == response.headers.get("x-litellm-call-id")
|
||||
|
||||
|
||||
class TestGuardrailBlockedResponsesShape:
|
||||
"""A pre_call block raises ModifyResponseException before any provider call.
|
||||
|
||||
|
|
|
|||
|
|
@ -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."}]}]
|
||||
),
|
||||
|
|
|
|||
|
|
@ -15778,7 +15778,7 @@ async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_coun
|
|||
warm_tokenizer("claude-fable-5")
|
||||
|
||||
response, took, lags = await timed_with_loop_lags(
|
||||
lambda: proxy_server_module.token_counter(TokenCountRequest(model="claude-fable-5", prompt=text * 100))
|
||||
lambda: proxy_server_module.count_request_tokens(TokenCountRequest(model="claude-fable-5", prompt=text * 100))
|
||||
)
|
||||
|
||||
assert response.total_tokens > 0
|
||||
|
|
@ -15819,7 +15819,7 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp
|
|||
|
||||
litellm.utils._select_custom_tokenizer_helper.cache_clear()
|
||||
try:
|
||||
response: Final = await proxy_server_module.token_counter(
|
||||
response: Final = await proxy_server_module.count_request_tokens(
|
||||
TokenCountRequest(model="self-hosted", prompt="count me off the loop")
|
||||
)
|
||||
assert loader_threads.get_nowait() is not loop_thread
|
||||
|
|
@ -15860,7 +15860,7 @@ async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revisi
|
|||
litellm.utils._select_custom_tokenizer_helper.cache_clear()
|
||||
try:
|
||||
responses: Final = [
|
||||
await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me once"))
|
||||
await proxy_server_module.count_request_tokens(TokenCountRequest(model="self-hosted", prompt="count me once"))
|
||||
for _ in range(3)
|
||||
]
|
||||
assert from_pretrained.call_args_list == [mock.call("my-org/tokenizer", revision="main", token=None)]
|
||||
|
|
@ -15868,8 +15868,8 @@ async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revisi
|
|||
assert len({response.total_tokens for response in responses}) == 1
|
||||
assert responses[0].total_tokens > 0
|
||||
|
||||
await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-pinned", prompt="count me once"))
|
||||
await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-private", prompt="count me once"))
|
||||
await proxy_server_module.count_request_tokens(TokenCountRequest(model="self-hosted-pinned", prompt="count me once"))
|
||||
await proxy_server_module.count_request_tokens(TokenCountRequest(model="self-hosted-private", prompt="count me once"))
|
||||
assert from_pretrained.call_args_list == [
|
||||
mock.call("my-org/tokenizer", revision="main", token=None),
|
||||
mock.call("my-org/tokenizer", revision="v2", token=None),
|
||||
|
|
|
|||
|
|
@ -38,7 +38,7 @@ load_dotenv()
|
|||
# this file is to test litellm/proxy
|
||||
|
||||
|
||||
from fastapi import Request
|
||||
from fastapi import Request, Response
|
||||
from fastapi.responses import JSONResponse
|
||||
from starlette.types import Message
|
||||
|
||||
|
|
@ -83,7 +83,6 @@ from litellm.proxy.proxy_server import (
|
|||
_provider_token_count,
|
||||
_try_provider_token_count,
|
||||
count_request_tokens,
|
||||
token_counter,
|
||||
)
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted
|
||||
|
|
@ -139,7 +138,7 @@ async def test_vLLM_token_counting():
|
|||
|
||||
setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
|
||||
|
||||
response = await token_counter(
|
||||
response = await count_request_tokens(
|
||||
request=TokenCountRequest(
|
||||
model="special-alias",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
|
|
@ -172,7 +171,7 @@ async def test_token_counting_model_not_in_model_list():
|
|||
|
||||
setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
|
||||
|
||||
response = await token_counter(
|
||||
response = await count_request_tokens(
|
||||
request=TokenCountRequest(
|
||||
model="special-alias",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
|
|
@ -209,8 +208,8 @@ async def test_asking_the_provider_endpoint_leaves_an_openrouter_deployments_cou
|
|||
model="or-gpt", messages=[{"role": "user", "content": "Count these tokens, ünïcödé 你好 🙂"}]
|
||||
)
|
||||
|
||||
via_provider_endpoint: Final = await token_counter(request=request, call_endpoint=True)
|
||||
counted_locally: Final = await token_counter(request=request, call_endpoint=False)
|
||||
via_provider_endpoint: Final = await count_request_tokens(request=request, call_endpoint=True)
|
||||
counted_locally: Final = await count_request_tokens(request=request, call_endpoint=False)
|
||||
|
||||
assert (via_provider_endpoint.total_tokens, via_provider_endpoint.model_used) == (
|
||||
counted_locally.total_tokens,
|
||||
|
|
@ -238,7 +237,7 @@ async def test_gpt_token_counting():
|
|||
|
||||
setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
|
||||
|
||||
response = await token_counter(
|
||||
response = await count_request_tokens(
|
||||
request=TokenCountRequest(
|
||||
model="gpt-4",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
|
|
@ -281,7 +280,7 @@ async def test_token_counter_leaves_the_routers_deployment_without_the_credentia
|
|||
)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
|
||||
|
||||
response = await token_counter(
|
||||
response = await count_request_tokens(
|
||||
request=TokenCountRequest(model="devstral", messages=[{"role": "user", "content": "hello"}])
|
||||
)
|
||||
|
||||
|
|
@ -312,10 +311,10 @@ async def test_token_counter_reads_an_oidc_api_key_only_when_it_calls_the_provid
|
|||
request = TokenCountRequest(model="claude", messages=[{"role": "user", "content": "hello"}])
|
||||
|
||||
if not call_endpoint:
|
||||
assert (await token_counter(request=request, call_endpoint=call_endpoint)).total_tokens > 0
|
||||
assert (await count_request_tokens(request=request, call_endpoint=call_endpoint)).total_tokens > 0
|
||||
return
|
||||
with pytest.raises(ProxyException) as excinfo:
|
||||
await token_counter(request=request, call_endpoint=call_endpoint)
|
||||
await count_request_tokens(request=request, call_endpoint=call_endpoint)
|
||||
assert excinfo.value.code == "401"
|
||||
assert str(tmp_path) not in excinfo.value.message
|
||||
|
||||
|
|
@ -343,7 +342,7 @@ async def test_token_counter_answers_503_when_the_metadata_server_behind_an_oidc
|
|||
500, text="metadata server unavailable"
|
||||
)
|
||||
with pytest.raises(ProxyException) as excinfo:
|
||||
await token_counter(
|
||||
await count_request_tokens(
|
||||
request=TokenCountRequest(model="claude", messages=[{"role": "user", "content": "hello"}]),
|
||||
call_endpoint=True,
|
||||
)
|
||||
|
|
@ -463,7 +462,7 @@ async def test_anthropic_messages_count_tokens_endpoint():
|
|||
|
||||
try:
|
||||
# Call the endpoint
|
||||
response = await count_tokens(mock_request, mock_user_api_key_dict)
|
||||
response = await count_tokens(mock_request, Response(), mock_user_api_key_dict)
|
||||
|
||||
# Verify response format matches Anthropic spec
|
||||
assert isinstance(response, dict)
|
||||
|
|
@ -545,7 +544,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model():
|
|||
|
||||
try:
|
||||
# Call the endpoint
|
||||
response = await count_tokens(mock_request, mock_user_api_key_dict)
|
||||
response = await count_tokens(mock_request, Response(), mock_user_api_key_dict)
|
||||
|
||||
# Verify response format matches Anthropic spec
|
||||
assert isinstance(response, dict)
|
||||
|
|
@ -586,7 +585,7 @@ async def test_internal_token_counter_anthropic_provider_detection():
|
|||
# Test with is_direct_request=False (simulating call from Anthropic endpoint)
|
||||
# The provider counter is engaged: its API error surfaces instead of a local count
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await token_counter(
|
||||
await count_request_tokens(
|
||||
request=TokenCountRequest(
|
||||
model="claude-test",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
|
|
@ -610,7 +609,7 @@ async def test_internal_token_counter_anthropic_provider_detection():
|
|||
setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
|
||||
|
||||
with pytest.raises(ProxyException) as missing_key:
|
||||
await token_counter(
|
||||
await count_request_tokens(
|
||||
request=TokenCountRequest(
|
||||
model="gpt-test",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
|
|
@ -653,7 +652,7 @@ async def test_anthropic_endpoint_error_handling():
|
|||
anthropic_endpoints.read_request_body = mock_read_request_body
|
||||
|
||||
try:
|
||||
response = await count_tokens(mock_request, mock_user_api_key_dict)
|
||||
response = await count_tokens(mock_request, Response(), mock_user_api_key_dict)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert json.loads(response.body) == {
|
||||
|
|
@ -947,7 +946,7 @@ async def test_vertex_ai_anthropic_token_counting():
|
|||
mock_handle_count_tokens.return_value = mock_token_response
|
||||
|
||||
# Test with messages format and call_endpoint=True
|
||||
response = await token_counter(
|
||||
response = await count_request_tokens(
|
||||
request=TokenCountRequest(
|
||||
model="vertex_ai/claude-3-5-sonnet-20241022",
|
||||
messages=[
|
||||
|
|
@ -1161,7 +1160,7 @@ async def test_token_counter_httpx_status_error_raises_proxy_exception():
|
|||
litellm.proxy.proxy_server.llm_router = mock_router
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await token_counter(
|
||||
await count_request_tokens(
|
||||
request=TokenCountRequest(
|
||||
model="claude-4-6-sonnet",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
|
|
@ -1227,7 +1226,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled():
|
|||
litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await token_counter(
|
||||
await count_request_tokens(
|
||||
request=TokenCountRequest(
|
||||
model="claude-bedrock",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
|
|
@ -1292,7 +1291,7 @@ async def test_proxy_token_counter_error_raises_when_enabled():
|
|||
|
||||
with patch("litellm.token_counter") as mock_local_counter:
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await token_counter(
|
||||
await count_request_tokens(
|
||||
request=TokenCountRequest(
|
||||
model="claude-bedrock",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
|
|
@ -1356,7 +1355,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format():
|
|||
proxy_server.count_request_tokens = mock_token_counter_error
|
||||
|
||||
try:
|
||||
response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict)
|
||||
response: Final[object] = await anthropic_count_tokens(mock_request, Response(), mock_user_api_key_dict)
|
||||
|
||||
assert isinstance(response, JSONResponse)
|
||||
assert response.status_code == 400
|
||||
|
|
@ -1415,7 +1414,7 @@ async def test_anthropic_endpoint_403_permission_error_format():
|
|||
proxy_server.count_request_tokens = mock_token_counter_error
|
||||
|
||||
try:
|
||||
response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict)
|
||||
response: Final[object] = await anthropic_count_tokens(mock_request, Response(), mock_user_api_key_dict)
|
||||
|
||||
assert isinstance(response, JSONResponse)
|
||||
assert response.status_code == 403
|
||||
|
|
@ -1474,7 +1473,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format():
|
|||
proxy_server.count_request_tokens = mock_token_counter_error
|
||||
|
||||
try:
|
||||
response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict)
|
||||
response: Final[object] = await anthropic_count_tokens(mock_request, Response(), mock_user_api_key_dict)
|
||||
|
||||
assert isinstance(response, JSONResponse)
|
||||
assert response.status_code == 429
|
||||
|
|
@ -1550,7 +1549,7 @@ async def test_anthropic_count_tokens_drops_a_malformed_tool_result_text_when_dr
|
|||
route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens").mock(
|
||||
return_value=httpx.Response(200, json={"input_tokens": 9})
|
||||
)
|
||||
response: Final = await anthropic_count_tokens(request, UserAPIKeyAuth())
|
||||
response: Final = await anthropic_count_tokens(request, Response(), UserAPIKeyAuth())
|
||||
|
||||
assert response["input_tokens"] == 9
|
||||
sent: Final = json.loads(route.calls.last.request.content)
|
||||
|
|
@ -1595,7 +1594,7 @@ async def _count_on_a_model_without_a_count_api(
|
|||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
monkeypatch.setattr(proxy_server, "proxy_config", proxy_server.ProxyConfig())
|
||||
with respx.mock:
|
||||
return await anthropic_count_tokens(request, UserAPIKeyAuth())
|
||||
return await anthropic_count_tokens(request, Response(), UserAPIKeyAuth())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1676,7 +1675,7 @@ async def test_anthropic_count_tokens_on_a_claude_deployment_drops_a_malformed_t
|
|||
route: Final = respx.post("https://api.anthropic.com/v1/messages/count_tokens").mock(
|
||||
return_value=httpx.Response(200, json={"input_tokens": 7})
|
||||
)
|
||||
response: Final = await anthropic_count_tokens(request, UserAPIKeyAuth())
|
||||
response: Final = await anthropic_count_tokens(request, Response(), UserAPIKeyAuth())
|
||||
|
||||
assert response["input_tokens"] == 7
|
||||
sent: Final = json.loads(route.calls.last.request.content)
|
||||
|
|
@ -2191,8 +2190,8 @@ async def test_framed_estimate_that_cannot_read_the_body_answers_the_provider_er
|
|||
)
|
||||
|
||||
with pytest.raises(ProxyException) as raised:
|
||||
await token_counter(request=TokenCountRequest(model="kimi", **unreadable_body), call_endpoint=True)
|
||||
readable: Final = await token_counter(
|
||||
await count_request_tokens(request=TokenCountRequest(model="kimi", **unreadable_body), call_endpoint=True)
|
||||
readable: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]),
|
||||
call_endpoint=True,
|
||||
)
|
||||
|
|
@ -2223,7 +2222,7 @@ async def test_framed_estimate_counts_a_json_schema_shorthand_like_its_explicit_
|
|||
)
|
||||
|
||||
async def count(schema: object) -> TokenCountResponse:
|
||||
return await token_counter(
|
||||
return await count_request_tokens(
|
||||
request=TokenCountRequest(
|
||||
model="kimi",
|
||||
messages=[{"role": "user", "content": "What is new?"}],
|
||||
|
|
@ -2277,8 +2276,8 @@ async def test_framed_estimate_counts_google_contents_like_the_same_conversation
|
|||
_framed_kimi(monkeypatch)
|
||||
google_request: Final = _google_count_tokens_request("kimi", {"contents": contents})
|
||||
|
||||
from_google: Final = await token_counter(request=google_request, call_endpoint=True)
|
||||
from_messages: Final = await token_counter(
|
||||
from_google: Final = await count_request_tokens(request=google_request, call_endpoint=True)
|
||||
from_messages: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model="kimi", messages=google_request.messages), call_endpoint=True
|
||||
)
|
||||
|
||||
|
|
@ -2309,7 +2308,7 @@ async def test_framed_deployment_whose_custom_tokenizer_cannot_load_answers_the_
|
|||
)
|
||||
|
||||
with pytest.raises(ProxyException) as raised:
|
||||
await token_counter(
|
||||
await count_request_tokens(
|
||||
request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]),
|
||||
call_endpoint=True,
|
||||
)
|
||||
|
|
@ -2361,7 +2360,7 @@ def _framed_kimi_with_custom_tokenizer(monkeypatch: pytest.MonkeyPatch, identifi
|
|||
|
||||
async def _count_status(model: str = "kimi") -> str:
|
||||
try:
|
||||
response: Final = await token_counter(
|
||||
response: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model=model, messages=[{"role": "user", "content": "What is new?"}]),
|
||||
call_endpoint=True,
|
||||
)
|
||||
|
|
@ -2625,8 +2624,8 @@ async def test_local_count_loads_the_tokenizer_again_after_a_failed_load(monkeyp
|
|||
request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}])
|
||||
|
||||
with pytest.raises(OSError, match="hub unreachable"):
|
||||
await token_counter(request=request)
|
||||
retried: Final = await token_counter(request=request)
|
||||
await count_request_tokens(request=request)
|
||||
retried: Final = await count_request_tokens(request=request)
|
||||
|
||||
assert retried.count_source == "local_estimate" and retried.total_tokens > 0
|
||||
|
||||
|
|
@ -2656,15 +2655,15 @@ async def test_framed_count_loads_the_tokenizer_again_only_after_the_window_of_a
|
|||
monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_MAX_RUNNING", 1)
|
||||
request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}])
|
||||
|
||||
first: Final = await token_counter(request=request, call_endpoint=True)
|
||||
second: Final = await token_counter(request=request, call_endpoint=True)
|
||||
first: Final = await count_request_tokens(request=request, call_endpoint=True)
|
||||
second: Final = await count_request_tokens(request=request, call_endpoint=True)
|
||||
|
||||
assert (first.tokenizer_type, second.tokenizer_type) == tokenizer_types
|
||||
|
||||
|
||||
async def _fork_guard_traceback_frames() -> int:
|
||||
with pytest.raises(ForkedAfterNativeRuntimeStarted) as raised:
|
||||
await token_counter(
|
||||
await count_request_tokens(
|
||||
request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]),
|
||||
call_endpoint=True,
|
||||
)
|
||||
|
|
@ -2703,7 +2702,7 @@ async def test_swallowed_estimate_error_is_logged_by_type_without_the_request_co
|
|||
verbose_proxy_logger.addHandler(caplog.handler)
|
||||
try:
|
||||
with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name), pytest.raises(ProxyException):
|
||||
await token_counter(
|
||||
await count_request_tokens(
|
||||
request=TokenCountRequest(
|
||||
model="kimi",
|
||||
messages=[
|
||||
|
|
@ -2748,9 +2747,9 @@ async def test_local_estimate_adds_the_configured_framing_once_per_request(monke
|
|||
)
|
||||
messages: Final = [{"role": "system", "content": "be brief"}, {"role": "user", "content": "hello there"}]
|
||||
|
||||
plain: Final = await token_counter(request=TokenCountRequest(model="plain", messages=messages))
|
||||
framed: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages))
|
||||
framed_prompt: Final = await token_counter(request=TokenCountRequest(model="framed", prompt="hello there"))
|
||||
plain: Final = await count_request_tokens(request=TokenCountRequest(model="plain", messages=messages))
|
||||
framed: Final = await count_request_tokens(request=TokenCountRequest(model="framed", messages=messages))
|
||||
framed_prompt: Final = await count_request_tokens(request=TokenCountRequest(model="framed", prompt="hello there"))
|
||||
|
||||
assert plain.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages)
|
||||
assert framed.total_tokens == plain.total_tokens + 100
|
||||
|
|
@ -2782,9 +2781,9 @@ async def test_framing_comes_from_the_row_the_deployment_provider_or_base_model_
|
|||
messages: Final = [{"role": "user", "content": "hello there"}]
|
||||
framing: Final = cost_map_entry(kimi)["token_count_framing"]
|
||||
|
||||
prefixed: Final = await token_counter(request=TokenCountRequest(model="prefixed", messages=messages))
|
||||
split: Final = await token_counter(request=TokenCountRequest(model="split", messages=messages))
|
||||
aliased: Final = await token_counter(request=TokenCountRequest(model="aliased", messages=messages))
|
||||
prefixed: Final = await count_request_tokens(request=TokenCountRequest(model="prefixed", messages=messages))
|
||||
split: Final = await count_request_tokens(request=TokenCountRequest(model="split", messages=messages))
|
||||
aliased: Final = await count_request_tokens(request=TokenCountRequest(model="aliased", messages=messages))
|
||||
|
||||
assert prefixed.total_tokens == litellm.token_counter(model="system.ai.kimi-k3", messages=messages) + framing
|
||||
assert split.total_tokens == prefixed.total_tokens
|
||||
|
|
@ -2964,7 +2963,7 @@ async def test_local_estimate_matches_the_billed_prompt_with_the_shipped_framing
|
|||
),
|
||||
)
|
||||
|
||||
counted: Final = await token_counter(
|
||||
counted: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(
|
||||
model="billed",
|
||||
messages=[{"role": "user", "content": "What is the weather in Paris?"}],
|
||||
|
|
@ -3363,7 +3362,9 @@ async def test_bedrock_kimi_k3_deployment_with_a_custom_tokenizer_keeps_counting
|
|||
monkeypatch.setattr(litellm.utils, "select_tokenizer", lambda model, custom_tokenizer=None: custom)
|
||||
messages: Final = [{"role": "user", "content": "What is new?"}]
|
||||
|
||||
counted: Final = await token_counter(request=TokenCountRequest(model="kimi", messages=messages), call_endpoint=True)
|
||||
counted: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model="kimi", messages=messages), call_endpoint=True
|
||||
)
|
||||
|
||||
assert counted.total_tokens == (
|
||||
litellm.token_counter(model=counted.model_used, messages=messages, custom_tokenizer=custom)
|
||||
|
|
@ -3398,7 +3399,7 @@ async def test_bedrock_kimi_k3_without_its_tokenizer_adds_the_template_offsets_t
|
|||
):
|
||||
_framed_kimi(monkeypatch)
|
||||
|
||||
counted: Final = await token_counter(
|
||||
counted: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model="kimi", messages=messages, tools=tools), call_endpoint=True
|
||||
)
|
||||
estimate: Final = litellm.token_counter(
|
||||
|
|
@ -3427,7 +3428,7 @@ async def test_slow_cost_map_lookup_does_not_stall_other_requests_on_the_worker(
|
|||
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None)
|
||||
monkeypatch.setattr(litellm, "get_model_info", slow_get_model_info)
|
||||
counting: Final = asyncio.create_task(
|
||||
token_counter(
|
||||
count_request_tokens(
|
||||
request=TokenCountRequest(model="slow-lookup-model", messages=[{"role": "user", "content": "hi"}])
|
||||
)
|
||||
)
|
||||
|
|
@ -3462,7 +3463,7 @@ async def test_uncountable_model_keeps_the_provider_error_when_local_counting_is
|
|||
monkeypatch.setattr(litellm, "disable_token_counter", True)
|
||||
|
||||
with pytest.raises(ProxyException) as raised:
|
||||
await token_counter(
|
||||
await count_request_tokens(
|
||||
request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True
|
||||
)
|
||||
|
||||
|
|
@ -3512,7 +3513,7 @@ async def test_configured_token_count_model_overrides_the_shipped_one_and_an_emp
|
|||
litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock")
|
||||
)
|
||||
|
||||
counted: Final = await token_counter(
|
||||
counted: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True
|
||||
)
|
||||
|
||||
|
|
@ -3550,7 +3551,7 @@ async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geog
|
|||
litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock")
|
||||
)
|
||||
|
||||
await token_counter(
|
||||
await count_request_tokens(
|
||||
request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True
|
||||
)
|
||||
|
||||
|
|
@ -3608,7 +3609,7 @@ async def test_region_lookup_for_a_shipped_count_model_leaves_the_event_loop_fre
|
|||
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only"))
|
||||
|
||||
async def count() -> TokenCountResponse:
|
||||
return await token_counter(
|
||||
return await count_request_tokens(
|
||||
request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]),
|
||||
call_endpoint=True,
|
||||
)
|
||||
|
|
@ -3636,7 +3637,7 @@ async def test_framing_still_applies_while_the_deployment_is_cooling_down(monkey
|
|||
)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router)
|
||||
messages: Final = [{"role": "user", "content": "hello there"}]
|
||||
warm: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages))
|
||||
warm: Final = await count_request_tokens(request=TokenCountRequest(model="framed", messages=messages))
|
||||
|
||||
router.cooldown_cache.add_deployment_to_cooldown(
|
||||
model_id="framed-deployment",
|
||||
|
|
@ -3644,7 +3645,7 @@ async def test_framing_still_applies_while_the_deployment_is_cooling_down(monkey
|
|||
exception_status=429,
|
||||
cooldown_time=60,
|
||||
)
|
||||
cooling: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages))
|
||||
cooling: Final = await count_request_tokens(request=TokenCountRequest(model="framed", messages=messages))
|
||||
|
||||
assert warm.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + 100
|
||||
assert cooling.total_tokens - 100 == litellm.token_counter(model="framed", messages=messages)
|
||||
|
|
@ -3738,7 +3739,7 @@ async def _count_through_the_messages_endpoint(
|
|||
return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)),
|
||||
),
|
||||
):
|
||||
response: Final = await anthropic_count_tokens(mock_request, MagicMock())
|
||||
response: Final = await anthropic_count_tokens(mock_request, Response(), MagicMock())
|
||||
return response, peer.request_bodies
|
||||
|
||||
|
||||
|
|
@ -4098,7 +4099,7 @@ async def _messages_endpoint_count(router: Router, request_body: Mapping[str, Js
|
|||
patch.object(proxy_server, "llm_router", router),
|
||||
patch.object(anthropic_endpoints, "read_request_body", AsyncMock(return_value=dict(request_body))),
|
||||
):
|
||||
return await anthropic_count_tokens(mock_request, MagicMock())
|
||||
return await anthropic_count_tokens(mock_request, Response(), MagicMock())
|
||||
|
||||
|
||||
class _MessagesCountReply(BaseModel):
|
||||
|
|
@ -4651,10 +4652,10 @@ async def test_call_endpoint_counts_a_prompt_like_the_same_text_sent_as_a_user_m
|
|||
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer))
|
||||
|
||||
with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
|
||||
from_prompt: Final = await token_counter(
|
||||
from_prompt: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model="gemini-counted", prompt=text), call_endpoint=True
|
||||
)
|
||||
from_message: Final = await token_counter(
|
||||
from_message: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": text}]),
|
||||
call_endpoint=True,
|
||||
)
|
||||
|
|
@ -4700,7 +4701,7 @@ async def test_call_endpoint_counts_a_prompt_on_a_vertex_ai_deployment(monkeypat
|
|||
with patch.object(
|
||||
VertexAIPartnerModelsTokenCounter, "handle_count_tokens_request", side_effect=_vertex_partner_word_count
|
||||
):
|
||||
from_prompt: Final = await token_counter(
|
||||
from_prompt: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model="vertex-counted", prompt=text), call_endpoint=True
|
||||
)
|
||||
|
||||
|
|
@ -4729,7 +4730,7 @@ async def test_prompt_next_to_a_conversation_answers_400_without_drop_params(
|
|||
patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client),
|
||||
pytest.raises(ProxyException) as rejected,
|
||||
):
|
||||
await token_counter(
|
||||
await count_request_tokens(
|
||||
request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint
|
||||
)
|
||||
|
||||
|
|
@ -4764,10 +4765,10 @@ async def test_drop_params_counts_a_conversation_without_the_prompt(
|
|||
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer))
|
||||
|
||||
with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
|
||||
with_prompt: Final = await token_counter(
|
||||
with_prompt: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint
|
||||
)
|
||||
conversation_only: Final = await token_counter(
|
||||
conversation_only: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model="gemini-counted", **conversation), call_endpoint=call_endpoint
|
||||
)
|
||||
|
||||
|
|
@ -4784,7 +4785,7 @@ async def test_call_endpoint_prompt_next_to_an_empty_conversation_counts_the_pro
|
|||
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer))
|
||||
|
||||
with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
|
||||
counted: Final = await token_counter(
|
||||
counted: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation), call_endpoint=True
|
||||
)
|
||||
|
||||
|
|
@ -4797,8 +4798,10 @@ async def test_local_count_of_a_prompt_next_to_empty_messages_counts_the_prompt(
|
|||
) -> None:
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router())
|
||||
|
||||
with_empty: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX, messages=[]))
|
||||
alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX))
|
||||
with_empty: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model="gemini-counted", prompt=_FOX, messages=[])
|
||||
)
|
||||
alone: Final = await count_request_tokens(request=TokenCountRequest(model="gemini-counted", prompt=_FOX))
|
||||
|
||||
assert with_empty.error is False
|
||||
assert with_empty.total_tokens == alone.total_tokens > 0
|
||||
|
|
@ -4826,13 +4829,13 @@ async def test_local_count_of_a_prompt_with_system_or_tools_counts_them_like_the
|
|||
) -> None:
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router())
|
||||
|
||||
from_prompt: Final = await token_counter(
|
||||
from_prompt: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation, **chat_fields)
|
||||
)
|
||||
from_message: Final = await token_counter(
|
||||
from_message: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": _FOX}], **chat_fields)
|
||||
)
|
||||
prompt_alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX))
|
||||
prompt_alone: Final = await count_request_tokens(request=TokenCountRequest(model="gemini-counted", prompt=_FOX))
|
||||
|
||||
assert from_prompt.error is False
|
||||
assert from_prompt.total_tokens == from_message.total_tokens > prompt_alone.total_tokens
|
||||
|
|
@ -4903,8 +4906,8 @@ async def test_call_endpoint_sends_a_prompt_to_the_provider_as_one_user_message(
|
|||
)
|
||||
|
||||
with _count_apis_answering(count_api):
|
||||
from_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True)
|
||||
from_message: Final = await token_counter(request=_FOX_AS_MESSAGE, call_endpoint=True)
|
||||
from_prompt: Final = await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=True)
|
||||
from_message: Final = await count_request_tokens(request=_FOX_AS_MESSAGE, call_endpoint=True)
|
||||
|
||||
assert (from_prompt.total_tokens, from_prompt.error) == (41, False)
|
||||
assert from_prompt.count_source != "local_estimate"
|
||||
|
|
@ -4925,9 +4928,9 @@ async def test_call_endpoint_answers_a_provider_error_on_a_prompt_like_on_the_sa
|
|||
|
||||
with _count_apis_answering(count_api):
|
||||
with pytest.raises(ProxyException) as from_prompt:
|
||||
await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True)
|
||||
await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=True)
|
||||
with pytest.raises(ProxyException) as from_message:
|
||||
await token_counter(request=_FOX_AS_MESSAGE, call_endpoint=True)
|
||||
await count_request_tokens(request=_FOX_AS_MESSAGE, call_endpoint=True)
|
||||
|
||||
assert from_prompt.value.code == str(status)
|
||||
assert f"scripted {status}" in from_prompt.value.message
|
||||
|
|
@ -4944,11 +4947,11 @@ async def test_framed_stand_in_counts_a_prompt_like_the_same_text_sent_as_a_user
|
|||
) -> None:
|
||||
_framed_kimi(monkeypatch)
|
||||
|
||||
from_prompt: Final = await token_counter(request=TokenCountRequest(model="kimi", prompt=_FOX), call_endpoint=True)
|
||||
from_message: Final = await token_counter(
|
||||
from_prompt: Final = await count_request_tokens(request=TokenCountRequest(model="kimi", prompt=_FOX), call_endpoint=True)
|
||||
from_message: Final = await count_request_tokens(
|
||||
request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": _FOX}]), call_endpoint=True
|
||||
)
|
||||
local_prompt: Final = await token_counter(request=TokenCountRequest(model="kimi", prompt=_FOX))
|
||||
local_prompt: Final = await count_request_tokens(request=TokenCountRequest(model="kimi", prompt=_FOX))
|
||||
|
||||
assert (from_prompt.count_source, from_prompt.error) == ("local_estimate", False)
|
||||
assert from_prompt.total_tokens == from_message.total_tokens > local_prompt.total_tokens
|
||||
|
|
@ -4970,9 +4973,9 @@ async def test_call_endpoint_counts_a_prompt_on_a_provider_without_a_count_api_a
|
|||
return httpx.Response(200, json={"input_tokens": 41}, request=request)
|
||||
|
||||
with _count_apis_answering(count_api):
|
||||
from_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True)
|
||||
from_message: Final = await token_counter(request=_FOX_AS_MESSAGE, call_endpoint=True)
|
||||
local_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=False)
|
||||
from_prompt: Final = await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=True)
|
||||
from_message: Final = await count_request_tokens(request=_FOX_AS_MESSAGE, call_endpoint=True)
|
||||
local_prompt: Final = await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=False)
|
||||
|
||||
assert sent == []
|
||||
assert (from_prompt.count_source, from_prompt.error) == ("local_estimate", False)
|
||||
|
|
@ -5008,8 +5011,8 @@ async def test_call_endpoint_counts_a_prompt_on_a_non_chat_deployment_as_raw_tex
|
|||
return httpx.Response(200, json={"input_tokens": 41}, request=request)
|
||||
|
||||
with _count_apis_answering(count_api):
|
||||
from_endpoint: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True)
|
||||
local: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=False)
|
||||
from_endpoint: Final = await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=True)
|
||||
local: Final = await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=False)
|
||||
|
||||
assert sent == []
|
||||
assert (from_endpoint.total_tokens, from_endpoint.error) == (local.total_tokens, False)
|
||||
|
|
@ -5044,7 +5047,7 @@ async def test_call_endpoint_sends_a_prompt_on_a_chat_mode_deployment_to_the_pro
|
|||
return httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 41}, request=request)
|
||||
|
||||
with _count_apis_answering(count_api):
|
||||
from_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True)
|
||||
from_prompt: Final = await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=True)
|
||||
|
||||
assert (from_prompt.total_tokens, from_prompt.error) == (41, False)
|
||||
assert len(sent) == 1
|
||||
|
|
@ -5061,8 +5064,8 @@ async def test_call_endpoint_counts_a_prompt_for_a_model_without_a_deployment_as
|
|||
)
|
||||
request: Final = TokenCountRequest(model="deepseek/deepseek-v4-flash", prompt=_FOX)
|
||||
|
||||
from_endpoint: Final = await token_counter(request=request, call_endpoint=True)
|
||||
local: Final = await token_counter(request=request, call_endpoint=False)
|
||||
from_endpoint: Final = await count_request_tokens(request=request, call_endpoint=True)
|
||||
local: Final = await count_request_tokens(request=request, call_endpoint=False)
|
||||
|
||||
assert (from_endpoint.total_tokens, from_endpoint.error) == (local.total_tokens, False)
|
||||
|
||||
|
|
@ -5102,7 +5105,7 @@ async def test_call_endpoint_answers_a_timeout_when_a_stalled_count_api_outlasts
|
|||
|
||||
with pytest.raises(ProxyException) as error:
|
||||
async with asyncio.timeout(10):
|
||||
await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True)
|
||||
await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=True)
|
||||
|
||||
assert error.value.code == "408"
|
||||
|
||||
|
|
@ -5144,7 +5147,7 @@ async def test_call_endpoint_names_the_status_of_a_count_api_error_with_an_empty
|
|||
|
||||
with _count_apis_answering(lambda request: httpx.Response(404, request=request)):
|
||||
with pytest.raises(ProxyException) as error:
|
||||
await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True)
|
||||
await count_request_tokens(request=_FOX_AS_PROMPT, call_endpoint=True)
|
||||
|
||||
assert error.value.code == "404"
|
||||
assert "404" in error.value.message
|
||||
|
|
@ -5494,8 +5497,8 @@ async def test_token_counter_counts_a_custom_azure_deployment_name_with_its_base
|
|||
)
|
||||
messages: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}]
|
||||
|
||||
named: Final = await token_counter(request=TokenCountRequest(model="named", messages=messages))
|
||||
served: Final = await token_counter(request=TokenCountRequest(model="served", messages=messages))
|
||||
named: Final = await count_request_tokens(request=TokenCountRequest(model="named", messages=messages))
|
||||
served: Final = await count_request_tokens(request=TokenCountRequest(model="served", messages=messages))
|
||||
|
||||
assert named.total_tokens == served.total_tokens
|
||||
|
||||
|
|
@ -5527,7 +5530,9 @@ def _utils_token_counter(
|
|||
),
|
||||
)
|
||||
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-test"))
|
||||
return TestClient(app).post("/utils/token_counter", json=dict(body), headers=dict(headers))
|
||||
return TestClient(app, raise_server_exceptions=False).post(
|
||||
"/utils/token_counter", json=dict(body), headers=dict(headers)
|
||||
)
|
||||
|
||||
|
||||
def test_utils_token_counter_logs_one_line_with_the_callers_call_id_model_and_count(
|
||||
|
|
@ -5609,6 +5614,38 @@ def test_utils_token_counter_failed_count_logs_its_status_in_place_of_a_count(
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("body", "headers", "status_code"),
|
||||
[
|
||||
(_COUNTED_HELLO, {}, 200),
|
||||
(_COUNTED_HELLO, {"x-litellm-call-id": "count-echo-utils"}, 200),
|
||||
({"model": "count-log-gpt"}, {}, 400),
|
||||
({**_COUNTED_HELLO, "prompt": "and a prompt"}, {"x-litellm-call-id": "count-echo-utils"}, 400),
|
||||
(
|
||||
{"model": "count-log-gpt", "messages": [{"role": "user", "content": "hi", "tool_calls": "x"}]},
|
||||
{"x-litellm-call-id": "count-echo-utils"},
|
||||
500,
|
||||
),
|
||||
],
|
||||
ids=["counted", "counted-with-client-id", "no-messages", "failed-count-with-client-id", "crash-with-client-id"],
|
||||
)
|
||||
def test_utils_token_counter_answers_the_call_id_its_count_line_names(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
token_count_log_lines: Callable[[], tuple[dict[str, object], ...]],
|
||||
body: Mapping[str, JsonValue],
|
||||
headers: Mapping[str, str],
|
||||
status_code: int,
|
||||
) -> None:
|
||||
response: Final = _utils_token_counter(monkeypatch, body, headers)
|
||||
|
||||
(line,) = token_count_log_lines()
|
||||
assert (response.status_code, response.headers.get("x-litellm-call-id")) == (
|
||||
status_code,
|
||||
line["litellm_call_id"],
|
||||
), response.text
|
||||
assert headers.get("x-litellm-call-id") in (None, line["litellm_call_id"])
|
||||
|
||||
|
||||
def test_utils_token_counter_json_line_keeps_its_own_fields_when_the_call_id_is_a_dict(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
|
|
|
|||
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -19885,13 +19885,13 @@ export interface paths {
|
|||
put?: never;
|
||||
/**
|
||||
* Token Counter
|
||||
* @description Args:
|
||||
* @description Returns:
|
||||
* TokenCountResponse
|
||||
*
|
||||
* Args:
|
||||
* request: TokenCountRequest
|
||||
* call_endpoint: bool - When set to "True" it will call the token counting endpoint - e.g Anthropic or Google AI Studio Token Counting APIs.
|
||||
* litellm_call_id: the x-litellm-call-id header, logged with the count so it joins the inference call's logs
|
||||
*
|
||||
* Returns:
|
||||
* TokenCountResponse
|
||||
*/
|
||||
post: operations["token_counter_utils_token_counter_post"];
|
||||
delete?: never;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue