scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-16 04:29:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 138

fix(count_tokens): count Anthropic-format requests for OpenAI the way /v1/messages inference converts them (#46759)

Browse source 
* fix(count_tokens): convert Anthropic-format messages for the OpenAI counter the way inference does

* test(count_tokens): check /v1/messages/count_tokens sends OpenAI the input /v1/messages inference sends

* fix(count_tokens): build the Anthropic count request without a cast and sync dashboard API types

* test(count_tokens): compare tools by value in the public api test

* fix(count_tokens): count Anthropic bodies the way /v1/messages inference sends them to OpenAI

Unwrap affinity-encoded reasoning, count the structured-output schema, flatten replayed web search results, join a list system prompt, and answer a message without a role with a 400 before calling OpenAI. messages_format is now internal to the Anthropic count route instead of a public /utils/token_counter field.

* test(count_tokens): expect the image detail the adapter now sends after merging main

* fix(count_tokens): type normalize_native_messages without the unimported Any

On Python 3.10 to 3.13 the annotation is evaluated at import, so the proxy failed to start with NameError

* test(count_tokens): expect the converted Responses input in the budgeted-key count integration test

---------

Co-authored-by: Mrinal Chanshetty <mrinal@berri.ai>
This commit is contained in:
Mateo Wang 2026-10-16 01:58:46 +00:00 • committed by GitHub
parent 627bcab42b
commit d25b58f744
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
21 changed files with 814 additions and 161 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/llms/anthropic/count_tokens/token_counter.py
View file

@ -10,7 +10,7 @@ from litellm._logging import verbose_logger
from litellm.exceptions import AuthenticationError
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.types.utils import LlmProviders, TokenCountResponse
from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse
# Global handler instance - reuse across all token counting requests
anthropic_count_tokens_handler: Final = AnthropicCountTokensHandler()
@ -34,6 +34,8 @@ class AnthropicTokenCounter(BaseTokenCounter):
request_model: str = "",
tools: list[dict[str, JsonValue]] | None = None,
system: JsonValue = None,
messages_format: TokenCountMessagesFormat = "anthropic",
output_format: JsonValue = None,
) -> TokenCountResponse | None:
"""
Count tokens using Anthropic's CountTokens API.

15
litellm/llms/anthropic/pass_through/messages/utils.py
View file

@ -129,6 +129,13 @@ def anthropic_system_to_openai_message(system: object) -> ChatCompletionSystemMe
return ChatCompletionSystemMessage(role="system", content=system)
def normalize_native_messages(messages: Sequence[object]) -> list[dict[str, object]]:
"""Apply the block rewrites every /v1/messages request gets before it reaches a provider."""
return flatten_unencrypted_web_search_results_in_anthropic_messages(
sanitize_tool_use_ids_in_anthropic_messages(strip_empty_content_blocks_from_anthropic_messages(messages))
)
def prepare_native_messages(
messages: list[dict[str, JsonValue]],
system: str | list[dict[str, JsonValue]] | None,
@ -140,13 +147,7 @@ def prepare_native_messages(
api_base: str | None = None,
presanitized: bool = False,
) -> tuple[list[dict[str, JsonValue]], str | list[dict[str, JsonValue]] | None]:
normalized: Final = (
messages
if presanitized
else flatten_unencrypted_web_search_results_in_anthropic_messages(
sanitize_tool_use_ids_in_anthropic_messages(strip_empty_content_blocks_from_anthropic_messages(messages))
)
)
normalized: Final = messages if presanitized else normalize_native_messages(messages)
return cast( # cast-ok: legacy normalizers and injection preserve the JSON message and system shapes
tuple[list[dict[str, JsonValue]], str | list[dict[str, JsonValue]] | None],
AnthropicCacheControlHook.maybe_inject_cache_control(

12
litellm/llms/anthropic/pass_through/responses_adapters/transformation.py
View file

@ -579,6 +579,14 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
result["summary"] = "detailed"
return result
@classmethod
def system_blocks_to_instructions(cls, system: object) -> str:
return "\n".join(
text
for block in cls._content_blocks(system)
if block.get("type") == "text" and isinstance(text := block.get("text"), str) and text
)
def translate_request(
self,
anthropic_request: AnthropicMessagesRequest,
@ -628,9 +636,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
if isinstance(system, str):
responses_kwargs["instructions"] = system
elif isinstance(system, list):
responses_kwargs["instructions"] = "\n".join(
filter(None, (b.get("text", "") for b in system if isinstance(b, dict) and b.get("type") == "text"))
)
responses_kwargs["instructions"] = self.system_blocks_to_instructions(system)
# max_tokens -> max_output_tokens
max_tokens: Final = anthropic_request.get("max_tokens")

4
litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py
View file

@ -12,7 +12,7 @@ from litellm.llms.azure_ai.anthropic.count_tokens.handler import (
AzureAIAnthropicCountTokensHandler,
)
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.types.utils import LlmProviders, TokenCountResponse
from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse
# Global handler instance - reuse across all token counting requests
azure_ai_anthropic_count_tokens_handler: Final = AzureAIAnthropicCountTokensHandler()
@ -36,6 +36,8 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
request_model: str = "",
tools: list[dict[str, JsonValue]] | None = None,
system: object | None = None,
messages_format: TokenCountMessagesFormat = "anthropic",
output_format: JsonValue = None,
) -> TokenCountResponse | None:
"""
Count tokens using Azure AI Anthropic's CountTokens API.

4
litellm/llms/base_llm/base_utils.py
View file

@ -16,7 +16,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountResponse
from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountMessagesFormat, TokenCountResponse
class BaseTokenCounter(ABC):
@ -30,6 +30,8 @@ class BaseTokenCounter(ABC):
request_model: str = "",
tools: list[dict[str, Any]] | None = None,
system: Any | None = None,
messages_format: TokenCountMessagesFormat = "openai",
output_format: JsonValue = None,
) -> TokenCountResponse | None:
pass

4
litellm/llms/bedrock/count_tokens/bedrock_token_counter.py
View file

@ -16,7 +16,7 @@ from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler
from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountResponse
from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountMessagesFormat, TokenCountResponse
RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api"
MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api"
@ -107,6 +107,8 @@ class BedrockTokenCounter(BaseTokenCounter):
request_model: str = "",
tools: Sequence[Mapping[str, object]] | None = None,
system: object | None = None,
messages_format: TokenCountMessagesFormat = "anthropic",
output_format: JsonValue = None,
) -> TokenCountResponse | None:
if not messages:
return None

6
litellm/llms/gemini/common_utils.py
View file

@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence
from typing import Any, Final
import httpx
from pydantic import TypeAdapter
from pydantic import JsonValue, TypeAdapter
import litellm
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
@ -14,7 +14,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import TokenCountResponse
from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
GEMINI_IMAGE_ASPECT_RATIOS: Final[dict[str, float]] = {
"1:1": 1 / 1,
@ -495,6 +495,8 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter):
request_model: str = "",
tools: list[dict[str, object]] | None = None,
system: object | None = None,
messages_format: TokenCountMessagesFormat = "openai",
output_format: JsonValue = None,
) -> TokenCountResponse | None:
import copy

6
litellm/llms/openai/responses/count_tokens/handler.py
View file

@ -5,7 +5,7 @@ Uses httpx for HTTP requests to OpenAI's /v1/responses/input_tokens endpoint.
"""
import json
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from typing import Final
import httpx
@ -31,8 +31,9 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig):
api_key: str,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
tools: list[dict[str, object]] | None = None,
tools: Sequence[Mapping[str, object]] | None = None,
instructions: str | None = None,
text: Mapping[str, object] | None = None,
) -> dict[str, object]:
"""
Handle a token counting request to OpenAI's Responses API.
@ -53,6 +54,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig):
input=input,
tools=tools,
instructions=instructions,
text=text,
)
endpoint_url: Final = self.get_openai_count_tokens_endpoint(api_base)

78
litellm/llms/openai/responses/count_tokens/token_counter.py
View file

@ -3,9 +3,16 @@ OpenAI Token Counter implementation using the Responses API /input_tokens endpoi
"""
import os
from collections.abc import Mapping, Sequence
from typing import Any, Final
from pydantic import JsonValue
from typing_extensions import assert_never
from litellm._logging import verbose_logger
from litellm.llms.anthropic.pass_through.responses_adapters.transformation import (
LiteLLMAnthropicToResponsesAPIAdapter,
)
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.llms.openai.common_utils import OpenAIError
from litellm.llms.openai.responses.count_tokens.handler import (
@ -13,8 +20,9 @@ from litellm.llms.openai.responses.count_tokens.handler import (
)
from litellm.llms.openai.responses.count_tokens.transformation import (
OpenAICountTokensConfig,
ResponsesCountRequestParts,
)
from litellm.types.utils import LlmProviders, TokenCountResponse
from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse
# Global handler instance - reuse across all token counting requests
openai_count_tokens_handler: Final = OpenAICountTokensHandler()
@ -29,6 +37,37 @@ class OpenAITokenCounter(BaseTokenCounter):
) -> bool:
return custom_llm_provider == LlmProviders.OPENAI.value
@staticmethod
def _system_instructions(system: object) -> str | None:
if system is None or isinstance(system, str):
return system
return LiteLLMAnthropicToResponsesAPIAdapter.system_blocks_to_instructions(system) or None
@staticmethod
def _responses_request_parts(
model_to_use: str,
messages: Sequence[Mapping[str, object]],
tools: Sequence[Mapping[str, object]] | None,
system: object,
messages_format: TokenCountMessagesFormat,
output_format: JsonValue,
) -> ResponsesCountRequestParts:
match messages_format:
case "anthropic":
return OpenAICountTokensConfig.anthropic_messages_to_responses_input(
model=model_to_use, messages=messages, system=system, tools=tools, output_format=output_format
)
case "openai":
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(
[dict(message) for message in messages]
)
fallback_instructions: Final = OpenAITokenCounter._system_instructions(system)
return OpenAICountTokensConfig.count_request_parts(
input_items, instructions if instructions is not None else fallback_instructions, tools
)
case _:
assert_never(messages_format)
async def count_tokens(
self,
model_to_use: str,
@ -38,6 +77,8 @@ class OpenAITokenCounter(BaseTokenCounter):
request_model: str = "",
tools: list[dict[str, object]] | None = None,
system: object = None,
messages_format: TokenCountMessagesFormat = "openai",
output_format: JsonValue = None,
) -> TokenCountResponse | None:
"""
Count tokens using OpenAI's Responses API /input_tokens endpoint.
@ -45,6 +86,22 @@ class OpenAITokenCounter(BaseTokenCounter):
if not messages:
return None
malformed_error: Final = (
OpenAICountTokensConfig.malformed_anthropic_messages_error(messages)
if messages_format == "anthropic"
else None
)
if malformed_error is not None:
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="openai_api",
error=True,
error_message=malformed_error,
status_code=400,
)
deployment = deployment or {}
litellm_params: Final = deployment.get("litellm_params", {})
@ -59,12 +116,14 @@ class OpenAITokenCounter(BaseTokenCounter):
api_base: Final = litellm_params.get("api_base")
# Convert chat messages to Responses API input format
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
# Use system param if instructions not extracted from messages
if instructions is None and system is not None:
instructions = system if isinstance(system, str) else str(system)
input_items, instructions, counted_tools, text = self._responses_request_parts(
model_to_use=model_to_use,
messages=messages,
tools=tools,
system=system,
messages_format=messages_format,
output_format=output_format,
)
# If no input items were produced (e.g., system-only messages), fall back to local counting
if not input_items:
@ -73,11 +132,12 @@ class OpenAITokenCounter(BaseTokenCounter):
try:
result: Final = await openai_count_tokens_handler.handle_count_tokens_request(
model=model_to_use,
input=input_items if input_items is not None else [],
input=input_items,
api_key=api_key,
api_base=api_base,
tools=tools,
tools=counted_tools,
instructions=instructions,
text=text,
)
if result is not None:

88
litellm/llms/openai/responses/count_tokens/transformation.py
View file

@ -5,10 +5,18 @@ This module handles the transformation of requests to OpenAI's /v1/responses/inp
"""
from collections.abc import Mapping, Sequence
from typing import Any, Final, Literal
from typing import Any, Final, Literal, TypeAlias
from pydantic import JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages
from litellm.llms.anthropic.pass_through.responses_adapters.transformation import (
LiteLLMAnthropicToResponsesAPIAdapter,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.anthropic import AnthropicMessagesRequest, AnthropicOutputSchema
class ResponsesInputTextPart(TypedDict):
type: ReadOnly[Literal["input_text"]]
@ -31,6 +39,30 @@ ResponsesInputPart = ResponsesInputTextPart | ResponsesInputImagePart | Response
ResponsesContentRole = Literal["user", "assistant"]
ResponsesCountRequestParts: TypeAlias = tuple[
tuple[Mapping[str, object], ...],
str | None,
tuple[Mapping[str, object], ...] | None,
Mapping[str, object] | None,
]
_RESPONSES_COUNT_REQUEST_PARTS: Final = TypeAdapter(ResponsesCountRequestParts)
_TRANSLATED_REQUEST: Final = TypeAdapter(Mapping[str, object])
class _AnthropicCountMessage(TypedDict):
role: ReadOnly[str]
_ANTHROPIC_COUNT_MESSAGES: Final = TypeAdapter(tuple[_AnthropicCountMessage, ...])
_ANTHROPIC_OUTPUT_SCHEMA: Final = TypeAdapter(AnthropicOutputSchema)
def _anthropic_output_schema(output_format: JsonValue) -> AnthropicOutputSchema | None:
try:
return _ANTHROPIC_OUTPUT_SCHEMA.validate_python(output_format)
except ValidationError:
return None
def _chat_image_block_to_responses_part(image_url: object) -> ResponsesInputImagePart | None:
url: Final = image_url.get("url") if isinstance(image_url, Mapping) else image_url
@ -118,8 +150,9 @@ class OpenAICountTokensConfig:
self,
model: str,
input: str | Sequence[object],
tools: list[dict[str, object]] | None = None,
tools: Sequence[Mapping[str, object]] | None = None,
instructions: str | None = None,
text: Mapping[str, object] | None = None,
) -> dict[str, object]:
"""
Transform request to OpenAI Responses API token counting format.
@ -135,7 +168,10 @@ class OpenAICountTokensConfig:
request["instructions"] = instructions
if tools is not None:
request["tools"] = self._transform_tools_for_responses_api(tools)
request["tools"] = self._transform_tools_for_responses_api([dict(tool) for tool in tools])
if text is not None:
request["text"] = text
return request
@ -194,7 +230,7 @@ class OpenAICountTokensConfig:
@staticmethod
def messages_to_responses_input(
messages: list[dict[str, Any]],
) -> tuple:
) -> tuple[list[dict[str, object]], str | None]:
"""
Convert standard chat messages format to OpenAI Responses API input format.
@ -257,3 +293,47 @@ class OpenAICountTokensConfig:
instructions: Final = "\n".join(instructions_parts) if instructions_parts else None
return input_items, instructions
@staticmethod
def malformed_anthropic_messages_error(messages: Sequence[Mapping[str, object]]) -> str | None:
try:
_ANTHROPIC_COUNT_MESSAGES.validate_python(messages)
except ValidationError as e:
return "; ".join(
f"messages.{'.'.join(str(part) for part in error['loc'])}: {error['msg']}" for error in e.errors()
)
return None
@staticmethod
def anthropic_messages_to_responses_input(
model: str,
messages: Sequence[Mapping[str, object]],
system: object,
tools: Sequence[Mapping[str, object]] | None,
output_format: JsonValue = None,
) -> ResponsesCountRequestParts:
"""Convert with the adapter /v1/messages inference uses, so the count sees the same input it bills."""
anthropic_request: Final = AnthropicMessagesRequest(
model=model,
messages=normalize_native_messages(messages),
system=system if isinstance(system, (str, list)) else None,
tools=[dict(tool) for tool in tools] if tools else None,
output_format=_anthropic_output_schema(output_format),
)
translated: Final = _TRANSLATED_REQUEST.validate_python(
LiteLLMAnthropicToResponsesAPIAdapter().translate_request(
anthropic_request, include_encrypted_reasoning=False
)
)
return OpenAICountTokensConfig.count_request_parts(
ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(translated.get("input")),
translated.get("instructions"),
translated.get("tools"),
translated.get("text"),
)
@staticmethod
def count_request_parts(
input_items: object, instructions: object, tools: object, text: object = None
) -> ResponsesCountRequestParts:
return _RESPONSES_COUNT_REQUEST_PARTS.validate_python((input_items, instructions, tools, text))

6
litellm/llms/vertex_ai/common_utils.py
View file

@ -8,7 +8,7 @@ from types import MappingProxyType
from typing import Any, Final, Literal, cast, get_type_hints
import httpx
from pydantic import TypeAdapter, ValidationError
from pydantic import JsonValue, TypeAdapter, ValidationError
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
@ -23,7 +23,7 @@ from litellm.types.llms.vertex_ai import (
PartType,
Schema,
)
from litellm.types.utils import TokenCountResponse
from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
from litellm.utils import supports_response_schema, supports_system_messages
GEMINI_VIDEO_METADATA_KEYS: Final = MappingProxyType(
@ -1278,6 +1278,8 @@ class VertexAITokenCounter(BaseTokenCounter):
request_model: str = "",
tools: list[dict[str, object]] | None = None,
system: object | None = None,
messages_format: TokenCountMessagesFormat = "openai",
output_format: JsonValue = None,
) -> TokenCountResponse | None:
import copy

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

@ -2,10 +2,12 @@
Unified /v1/messages endpoint - (Anthropic Spec)
"""
from collections.abc import Mapping
from typing import Final
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse
from pydantic import JsonValue
import litellm
from litellm.anthropic_interface.exceptions import (
@ -289,6 +291,14 @@ async def anthropic_response(
)
def _requested_output_format(data: Mapping[str, JsonValue]) -> JsonValue:
output_format: Final = data.get("output_format")
output_config: Final = data.get("output_config")
if isinstance(output_format, dict) or not isinstance(output_config, dict):
return output_format
return output_config.get("format")
@router.post(
"/v1/messages/count_tokens",
tags=["[beta] Anthropic Messages Token Counting"],
@ -318,7 +328,7 @@ async def count_tokens(
Returns: {"input_tokens": <number>}
"""
from litellm.proxy.proxy_server import token_counter as internal_token_counter
from litellm.proxy.proxy_server import count_request_tokens
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
try:
@ -345,10 +355,11 @@ async def count_tokens(
system=data.get("system"),
)
# Call the internal token counter function with direct request flag set to False
token_response: Final = await internal_token_counter(
token_response: Final = await count_request_tokens(
request=token_request,
call_endpoint=True,
messages_format="anthropic",
output_format=_requested_output_format(data),
)
_token_response_dict: dict = {}
if isinstance(token_response, TokenCountResponse):

16
litellm/proxy/proxy_server.py
View file

@ -173,6 +173,7 @@ from litellm.types.utils import (
ModelResponseStream,
StreamingChoices,
TextCompletionResponse,
TokenCountMessagesFormat,
TokenCountResponse,
echoed_cost_map_pricing_fields,
is_server_derived_pricing_key,
@ -14415,6 +14416,8 @@ async def _try_provider_token_count(
request_model: str,
tools: list | None = None,
system: str | None = None,
messages_format: TokenCountMessagesFormat = "openai",
output_format: JsonValue = None,
) -> Optional["TokenCountResponse"]:
"""Attempt provider-specific token counting. Returns result on success, None to fall through to local counting."""
if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider):
@ -14428,6 +14431,8 @@ async def _try_provider_token_count(
request_model=request_model,
tools=tools,
system=system,
messages_format=messages_format,
output_format=output_format,
)
except httpx.HTTPStatusError as e:
error_message: Final = getattr(e, "message", None) or str(e)
@ -14470,8 +14475,15 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
Returns:
TokenCountResponse
"""
global llm_router
return await count_request_tokens(request=request, call_endpoint=call_endpoint)
async def count_request_tokens(
request: TokenCountRequest,
call_endpoint: bool = False,
messages_format: TokenCountMessagesFormat = "openai",
output_format: JsonValue = None,
) -> TokenCountResponse:
prompt: Final = request.prompt
messages: Final = request.messages
contents: Final = request.contents
@ -14530,6 +14542,8 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
request_model=request.model,
tools=tools,
system=system,
messages_format=messages_format,
output_format=output_format,
)
if result is not None:
return result

4
litellm/types/utils.py
View file

@ -9,6 +9,7 @@ from typing import (
Any,
Final,
Literal,
TypeAlias,
get_args,
)
@ -4471,6 +4472,9 @@ class LiteLLMLoggingBaseClass:
pass
TokenCountMessagesFormat: TypeAlias = Literal["openai", "anthropic"]
class TokenCountResponse(LiteLLMPydanticObjectBase):
total_tokens: int
request_model: str

192
tests/integration/providers/test_openai_count_tokens_tools_wire.py
View file

@ -1,5 +1,5 @@
import json
from collections.abc import Iterator
from collections.abc import Iterator, Mapping
from typing import Final
import pytest
@ -36,6 +36,146 @@ def _count_response(_request: Request) -> Reply:
return Reply(body=json.dumps({"input_tokens": _COUNT}).encode())
_PDF_BASE64: Final = "JVBERi0xLjQK"
_PNG_BASE64: Final = "aW1hZ2U="
_LOOKUP_TOOL: Final[dict[str, JsonValue]] = {
"name": "lookup",
"description": "Find a record",
"input_schema": {"type": "object", "properties": {"query": {"type": "string"}}},
}
_RESPONSES_REPLY: Final[Mapping[str, JsonValue]] = {
"id": "resp_count_parity",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-4o-mini",
"output": [
{
"type": "message",
"id": "msg_count_parity",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "OK", "annotations": []}],
}
],
"usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11},
}
_WRAPPED_REASONING: Final = "litellm_enc:bW9kZWxfaWQ6ZGVwbG95bWVudC0x;gAAAAoriginal"
_ANSWER_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}}
_ANTHROPIC_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = {
"structured-output": {
"output_config": {"format": {"type": "json_schema", "schema": _ANSWER_SCHEMA}},
"messages": [{"role": "user", "content": "Answer in JSON."}],
},
"affinity-wrapped-reasoning": {
"messages": [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": [
{"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{_WRAPPED_REASONING}"},
{"type": "text", "text": "OK"},
],
},
{"role": "user", "content": "again"},
],
},
"base64-pdf": {
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": "Summarize the attached contract."},
{
"type": "document",
"source": {"type": "base64", "media_type": "application/pdf", "data": _PDF_BASE64},
},
],
}
],
},
"base64-image": {
"system": [{"type": "text", "text": "Describe images briefly."}],
"messages": [
{
"role": "user",
"content": [
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": _PNG_BASE64}},
{"type": "text", "text": "What is this?"},
],
}
],
},
"replayed-web-search": {
"messages": [
{"role": "user", "content": "What is the newest LiteLLM release?"},
{
"role": "assistant",
"content": [
{
"type": "server_tool_use",
"id": "srvtoolu_01",
"name": "web_search",
"input": {"query": "newest LiteLLM release"},
},
{
"type": "web_search_tool_result",
"tool_use_id": "srvtoolu_01",
"content": [
{
"type": "web_search_result",
"url": "https://docs.litellm.ai/release_notes",
"title": "Release notes",
"encrypted_content": "",
"snippet": "LiteLLM ships a weekly release candidate and a stable line.",
}
],
},
{"type": "text", "text": "The newest release is in the release notes."},
],
},
{"role": "user", "content": "Thanks."},
],
},
"tool-turns": {
"system": "Use the lookup tool.",
"tools": [_LOOKUP_TOOL],
"messages": [
{"role": "user", "content": "Find record 14."},
{
"role": "assistant",
"content": [
{"type": "text", "text": "Looking it up."},
{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "14"}},
],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}],
},
],
},
}
def _responses_peer(request: Request) -> Reply:
if request.target == "/v1/responses/input_tokens":
return _count_response(request)
if request.target == "/v1/responses":
return Reply(body=json.dumps(_RESPONSES_REPLY).encode())
return Reply(status=404, body=json.dumps({"error": request.target}).encode())
def _body_for(wire_requests: tuple[Request, ...], target: str) -> dict[str, JsonValue]:
matching: Final = tuple(request for request in wire_requests if request.target == target)
assert len(matching) == 1, [request.target for request in wire_requests]
return _JSON_OBJECT.validate_json(matching[0].body)
def _counted_fields(body: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
return {key: body[key] for key in ("input", "instructions", "tools", "text") if body.get(key) is not None}
def _count_body(wire: Wire) -> dict[str, JsonValue]:
requests: Final = wire.drain()
assert len(requests) == 1, requests
@ -159,3 +299,53 @@ def test_utils_token_counter_keeps_chat_tool_conversion(counting_proxy: OwnedPro
"parameters": _TOOL_SCHEMA,
}
]
@pytest.mark.parametrize("case", sorted(_ANTHROPIC_BODIES))
def test_messages_count_tokens_sends_the_input_v1_messages_inference_sends(
counting_proxy: OwnedProxy, case: str
) -> None:
gateway: Final = counting_proxy.gateway
with wire_server(_responses_peer) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire.url)
body: Final = {"model": model, **_ANTHROPIC_BODIES[case]}
generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16})
counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body)
assert generated.status_code == 200, generated.text
assert counted.status_code == 200, counted.text
assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT
requests: Final = wire.drain()
inference_body: Final = _body_for(requests, "/v1/responses")
count_body: Final = _body_for(requests, "/v1/responses/input_tokens")
assert _counted_fields(count_body) == _counted_fields(inference_body)
assert count_body["model"] == inference_body["model"]
def test_messages_count_tokens_without_a_role_answers_400_before_calling_openai(counting_proxy: OwnedProxy) -> None:
gateway: Final = counting_proxy.gateway
with wire_server(_count_response) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire.url)
response: Final = gateway.request(
"POST",
"/v1/messages/count_tokens",
{"model": model, "messages": [{"role": "user", "content": "hello"}, {"content": "no role"}]},
)
assert response.status_code == 400, response.text
assert "messages.1.role: Field required" in response.text
assert wire.drain() == ()
def test_google_count_tokens_sends_the_system_instruction_text(counting_proxy: OwnedProxy) -> None:
gateway: Final = counting_proxy.gateway
with wire_server(_count_response) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire.url)
response: Final = gateway.request(
"POST",
f"/v1beta/models/{model}:countTokens",
{
"systemInstruction": {"parts": [{"text": "Be brief."}, {"text": "Answer in English."}]},
"contents": [{"role": "user", "parts": [{"text": "hello"}]}],
},
)
assert response.status_code == 200, response.text
assert _count_body(wire)["instructions"] == "Be brief.\nAnswer in English."

7
tests/integration/spend/test_cache_and_quota.py
View file

@ -369,7 +369,12 @@ def test_repeated_count_tokens_on_budgeted_key_does_not_reserve_budget_or_block_
{
"path": f"/{counting_upstream.scenario_id}/v1/responses/input_tokens",
"authorization": "Bearer integration-provider-key",
"body": {"model": "gpt-4o-mini", "input": [{"role": "user", "content": "hello!!!"}]},
"body": {
"model": "gpt-4o-mini",
"input": [
{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hello!!!"}]}
],
},
"method": "POST",
"api_key": "",
}

207
tests/unit/llms/openai/responses/count_tokens/test_token_counter.py
View file

@ -5,11 +5,12 @@ from typing import Final
import httpx
import pytest
import respx
from pydantic import ValidationError
from pydantic import JsonValue, ValidationError
import litellm
from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter
from litellm.types.utils import TokenCountResponse
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
COUNT_URL: Final = "https://api.openai.com/v1/responses/input_tokens"
@ -163,6 +164,208 @@ async def test_upstream_error_status_becomes_an_error_response_with_that_status_
)
async def _sent_count_body(
monkeypatch: pytest.MonkeyPatch,
messages: list[dict[str, object]],
system: object,
messages_format: TokenCountMessagesFormat,
output_format: JsonValue = None,
) -> object:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
with respx.mock:
route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1}))
await OpenAITokenCounter().count_tokens(
model_to_use="gpt-4o",
messages=messages,
contents=None,
deployment={"litellm_params": {"api_key": "sk-test"}},
request_model="openai/gpt-4o",
system=system,
messages_format=messages_format,
output_format=output_format,
)
return json.loads(route.calls.last.request.content)
@pytest.mark.asyncio
async def test_anthropic_format_sends_the_document_image_and_tool_turns_to_the_count_endpoint(
monkeypatch: pytest.MonkeyPatch,
) -> None:
messages: Final[list[dict[str, object]]] = [
{
"role": "user",
"content": [
{"type": "text", "text": "Summarize."},
{"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBE"}},
{"type": "image", "source": {"type": "url", "url": "https://example.com/page.png"}},
],
},
{"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]},
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "found"}]},
]
sent: Final = await _sent_count_body(
monkeypatch, messages, [{"type": "text", "text": "Be brief."}], messages_format="anthropic"
)
assert sent == {
"model": "gpt-4o",
"input": [
{
"type": "message",
"role": "user",
"content": [
{"type": "input_text", "text": "Summarize."},
{"type": "input_file", "filename": "document.pdf", "file_data": "data:application/pdf;base64,JVBE"},
{"type": "input_image", "image_url": "https://example.com/page.png", "detail": "auto"},
],
},
{"type": "function_call", "call_id": "toolu_01", "name": "lookup", "arguments": "{}"},
{"type": "function_call_output", "call_id": "toolu_01", "output": "found"},
],
"instructions": "Be brief.",
}
@pytest.mark.asyncio
async def test_openai_format_keeps_the_chat_conversion(monkeypatch: pytest.MonkeyPatch) -> None:
messages: Final[list[dict[str, object]]] = [
{"role": "system", "content": "Be brief."},
{"role": "user", "content": [{"type": "text", "text": "Summarize."}]},
{"role": "tool", "tool_call_id": "call_1", "content": "found"},
]
sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="openai")
assert sent == {
"model": "gpt-4o",
"input": [
{"role": "user", "content": "Summarize."},
{"type": "function_call_output", "call_id": "call_1", "output": "found"},
],
"instructions": "Be brief.",
}
@pytest.mark.asyncio
async def test_anthropic_format_sends_affinity_wrapped_reasoning_unwrapped_as_inference_does(
monkeypatch: pytest.MonkeyPatch,
) -> None:
wrapped: Final = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("gAAAAoriginal", "deployment-1")
messages: Final[list[dict[str, object]]] = [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": [
{"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"},
{"type": "text", "text": "OK"},
],
},
{"role": "user", "content": "again"},
]
sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic")
assert sent == {
"model": "gpt-4o",
"input": [
{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]},
{"type": "reasoning", "summary": [], "encrypted_content": "gAAAAoriginal"},
{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "OK"}]},
{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "again"}]},
],
}
@pytest.mark.asyncio
async def test_anthropic_format_counts_the_structured_output_schema(monkeypatch: pytest.MonkeyPatch) -> None:
schema: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}}
sent: Final = await _sent_count_body(
monkeypatch,
[{"role": "user", "content": "Answer in JSON."}],
system=None,
messages_format="anthropic",
output_format={"type": "json_schema", "schema": schema},
)
assert sent == {
"model": "gpt-4o",
"input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Answer in JSON."}]}],
"text": {"format": {"type": "json_schema", "name": "structured_output", "schema": schema, "strict": False}},
}
@pytest.mark.asyncio
async def test_anthropic_format_sends_a_replayed_web_search_result_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None:
snippet: Final = "LiteLLM ships a weekly release candidate and a stable line."
messages: Final[list[dict[str, object]]] = [
{"role": "user", "content": "What is new?"},
{
"role": "assistant",
"content": [
{"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}},
{
"type": "web_search_tool_result",
"tool_use_id": "srvtoolu_01",
"content": [
{
"type": "web_search_result",
"url": "https://docs.litellm.ai/release_notes",
"title": "Release notes",
"encrypted_content": "",
"snippet": snippet,
}
],
},
],
},
{"role": "user", "content": "Thanks."},
]
sent: Final = json.dumps(await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic"))
assert snippet in sent
assert "https://docs.litellm.ai/release_notes" in sent
@pytest.mark.asyncio
async def test_anthropic_format_message_without_a_role_answers_400_without_calling_openai(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
with respx.mock:
route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1}))
result: Final = await OpenAITokenCounter().count_tokens(
model_to_use="gpt-4o",
messages=[{"role": "user", "content": "hi"}, {"content": "no role"}],
contents=None,
deployment={"litellm_params": {"api_key": "sk-test"}},
request_model="openai/gpt-4o",
messages_format="anthropic",
)
assert route.called is False
assert result is not None
assert (result.error, result.status_code, result.error_message) == (True, 400, "messages.1.role: Field required")
@pytest.mark.asyncio
async def test_openai_format_sends_a_list_system_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None:
sent: Final = await _sent_count_body(
monkeypatch,
[{"role": "user", "content": "hi"}],
system=[{"type": "text", "text": "Be brief."}, {"type": "text", "text": "Answer in English."}],
messages_format="openai",
)
assert sent == {
"model": "gpt-4o",
"input": [{"role": "user", "content": "hi"}],
"instructions": "Be brief.\nAnswer in English.",
}
@pytest.mark.asyncio
async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message(
monkeypatch: pytest.MonkeyPatch,

108
tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py
View file

@ -2,6 +2,9 @@ from typing import Final
import pytest
from litellm.llms.anthropic.pass_through.responses_adapters.transformation import (
LiteLLMAnthropicToResponsesAPIAdapter,
)
from litellm.llms.openai.responses.count_tokens.transformation import (
OpenAICountTokensConfig,
)
@ -484,3 +487,108 @@ def test_get_required_headers():
assert headers["Authorization"] == "Bearer sk-test-key"
assert headers["Content-Type"] == "application/json"
PDF_BASE64: Final = "JVBERi0xLjcK"
PNG_BASE64: Final = "iVBORw0KGgo="
ANTHROPIC_AGENT_TURN: Final[list[dict[str, object]]] = [
{
"role": "user",
"content": [
{"type": "text", "text": "Summarize the indemnity terms."},
{
"type": "document",
"title": "contract.pdf",
"source": {"type": "base64", "media_type": "application/pdf", "data": PDF_BASE64},
},
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PNG_BASE64}},
],
},
{
"role": "assistant",
"content": [
{"type": "text", "text": "Let me pull clause 14."},
{"type": "tool_use", "id": "toolu_01", "name": "lookup_clause", "input": {"id": "14"}},
],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}],
},
]
ANTHROPIC_TOOLS: Final[list[dict[str, object]]] = [
{
"name": "lookup_clause",
"description": "Fetch a clause by id",
"input_schema": {"type": "object", "properties": {"id": {"type": "string"}}},
}
]
@pytest.mark.parametrize(
"system",
[None, "You review contracts.", [{"type": "text", "text": "You review contracts."}]],
)
def test_anthropic_messages_convert_exactly_as_v1_messages_inference_converts_them(system: object) -> None:
inference_request: Final[dict[str, object]] = {
"model": "gpt-4o",
"messages": ANTHROPIC_AGENT_TURN,
"tools": ANTHROPIC_TOOLS,
**({"system": system} if system is not None else {}),
}
inference: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_request(
inference_request, include_encrypted_reasoning=False
)
counted: Final = OpenAICountTokensConfig.anthropic_messages_to_responses_input(
model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=system, tools=ANTHROPIC_TOOLS
)
assert counted == (
tuple(inference["input"]),
inference.get("instructions"),
tuple(inference["tools"]),
inference.get("text"),
)
def test_anthropic_messages_keep_every_block_the_chat_converter_drops() -> None:
input_items, _, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input(
model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=None, tools=None
)
assert input_items == (
{
"type": "message",
"role": "user",
"content": [
{"type": "input_text", "text": "Summarize the indemnity terms."},
{
"type": "input_file",
"filename": "contract.pdf",
"file_data": f"data:application/pdf;base64,{PDF_BASE64}",
},
{"type": "input_image", "image_url": f"data:image/png;base64,{PNG_BASE64}", "detail": "auto"},
],
},
{"type": "function_call", "call_id": "toolu_01", "name": "lookup_clause", "arguments": '{"id": "14"}'},
{
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": "Let me pull clause 14."}],
},
{"type": "function_call_output", "call_id": "toolu_01", "output": "14. Indemnification."},
)
def test_anthropic_list_system_becomes_joined_instruction_text() -> None:
_, instructions, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input(
model="gpt-4o",
messages=[{"role": "user", "content": "hi"}],
system=[{"type": "text", "text": "First rule."}, {"type": "text", "text": "Second rule."}],
tools=None,
)
assert instructions == "First rule.\nSecond rule."

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

@ -278,13 +278,17 @@ class TestHttpExceptionDictDetail:
request.headers = {}
with (
patch.object(ep, "read_request_body", new=AsyncMock(return_value={})), # test-quality-ok: endpoint reads the body via a module function; no injection seam
patch.object(
ep, "read_request_body", new=AsyncMock(return_value={})
), # test-quality-ok: endpoint reads the body via a module function; no injection seam
patch.object( # test-quality-ok: the guardrail raise happens deep inside this call; the test targets the endpoint's except block
ep.ProxyBaseLLMRequestProcessing,
"base_process_llm_request",
new=AsyncMock(side_effect=exc),
),
patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam
patch.object(
proxy_server, "proxy_logging_obj"
) as mock_logging, # test-quality-ok: module global imported at call time; no injection seam
):
mock_logging.post_call_failure_hook = AsyncMock()
response = await ep.anthropic_response(
@ -375,9 +379,15 @@ class TestErrorLogCarriesCallId:
request.headers = {}
with (
patch.object(ep, "read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})), # test-quality-ok: endpoint reads the body via a module function; no injection seam
patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process), # test-quality-ok: the provider failure happens inside this call; the test targets the endpoint's except block
patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam
patch.object(
ep, "read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})
), # test-quality-ok: endpoint reads the body via a module function; no injection seam
patch.object(
ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process
), # test-quality-ok: the provider failure happens inside this call; the test targets the endpoint's except block
patch.object(
proxy_server, "proxy_logging_obj"
) as mock_logging, # test-quality-ok: module global imported at call time; no injection seam
caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"),
):
mock_logging.post_call_failure_hook = AsyncMock()
@ -408,9 +418,15 @@ class TestErrorLogCarriesCallId:
request.headers = {}
with (
patch.object(ep, "read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})), # test-quality-ok: endpoint reads the body via a module function; no injection seam
patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process), # test-quality-ok: the proxy shaped failure happens inside this call; the test targets the endpoint's except block
patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam
patch.object(
ep, "read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})
), # test-quality-ok: endpoint reads the body via a module function; no injection seam
patch.object(
ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process
), # test-quality-ok: the proxy shaped failure happens inside this call; the test targets the endpoint's except block
patch.object(
proxy_server, "proxy_logging_obj"
) as mock_logging, # test-quality-ok: module global imported at call time; no injection seam
):
mock_logging.post_call_failure_hook = AsyncMock()
response = await ep.anthropic_response(
@ -440,7 +456,9 @@ class TestErrorLogCarriesCallId:
"read_request_body",
new=AsyncMock(return_value={"model": "claude-sonnet", "messages": [{"role": "user", "content": "hi"}]}),
),
patch.object(proxy_server, "token_counter", new=AsyncMock(side_effect=RuntimeError("tokenizer down"))), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block
patch.object(
proxy_server, "count_request_tokens", new=AsyncMock(side_effect=RuntimeError("tokenizer down"))
), # 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"),
pytest.raises(HTTPException) as raised,
):

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

@ -30,7 +30,6 @@ from litellm.proxy.proxy_server import token_counter
from litellm.types.utils import TokenCountResponse
@pytest.mark.asyncio
async def test_vLLM_token_counting():
"""
@ -63,9 +62,7 @@ async def test_vLLM_token_counting():
print("response: ", response)
assert (
response.tokenizer_type == "openai_tokenizer"
) # SHOULD use the default tokenizer
assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the default tokenizer
assert response.model_used == "wolfram/miquliz-120b-v2.0"
@ -98,9 +95,7 @@ async def test_token_counting_model_not_in_model_list():
print("response: ", response)
assert (
response.tokenizer_type == "openai_tokenizer"
) # SHOULD use the OpenAI tokenizer
assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the OpenAI tokenizer
assert response.model_used == "special-alias"
@ -133,9 +128,7 @@ async def test_gpt_token_counting():
print("response: ", response)
assert (
response.tokenizer_type == "openai_tokenizer"
) # SHOULD use the OpenAI tokenizer
assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the OpenAI tokenizer
assert response.request_model == "gpt-4"
@ -173,10 +166,8 @@ async def test_anthropic_messages_count_tokens_endpoint():
anthropic_endpoints.read_request_body = mock_read_request_body
# Mock the internal token_counter function to return a controlled response
async def mock_token_counter(request, call_endpoint=False):
assert (
call_endpoint == True
), "Should be called with call_endpoint=True for Anthropic endpoint"
async def mock_token_counter(request, call_endpoint=False, messages_format="openai", output_format=None):
assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint"
assert request.model == "claude-3-sonnet-20240229"
assert request.messages == [{"role": "user", "content": "Hello Claude!"}]
@ -192,8 +183,8 @@ async def test_anthropic_messages_count_tokens_endpoint():
# Patch the imported token_counter function from proxy_server
import litellm.proxy.proxy_server as proxy_server
original_token_counter = proxy_server.token_counter
proxy_server.token_counter = mock_token_counter
original_token_counter = proxy_server.count_request_tokens
proxy_server.count_request_tokens = mock_token_counter
try:
# Call the endpoint
@ -210,7 +201,7 @@ async def test_anthropic_messages_count_tokens_endpoint():
finally:
# Restore original functions
anthropic_endpoints.read_request_body = original_read_request_body
proxy_server.token_counter = original_token_counter
proxy_server.count_request_tokens = original_token_counter
@pytest.mark.asyncio
@ -247,10 +238,8 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model():
anthropic_endpoints.read_request_body = mock_read_request_body
# Mock the internal token_counter function to return a controlled response
async def mock_token_counter(request, call_endpoint=True):
assert (
call_endpoint == True
), "Should be called with call_endpoint=True for Anthropic endpoint"
async def mock_token_counter(request, call_endpoint=True, messages_format="openai", output_format=None):
assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint"
assert request.model == "gpt-4"
assert request.messages == [{"role": "user", "content": "Hello GPT!"}]
@ -266,8 +255,8 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model():
# Patch the imported token_counter function from proxy_server
import litellm.proxy.proxy_server as proxy_server
original_token_counter = proxy_server.token_counter
proxy_server.token_counter = mock_token_counter
original_token_counter = proxy_server.count_request_tokens
proxy_server.count_request_tokens = mock_token_counter
try:
# Call the endpoint
@ -284,7 +273,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model():
finally:
# Restore original functions
anthropic_endpoints.read_request_body = original_read_request_body
proxy_server.token_counter = original_token_counter
proxy_server.count_request_tokens = original_token_counter
@pytest.mark.asyncio
@ -407,9 +396,7 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter():
# Mock the global handler instance in token_counter module
mock_handler = MagicMock()
mock_handler.handle_count_tokens_request = AsyncMock(
return_value={"input_tokens": 42}
)
mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42})
with patch(
"litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler",
@ -466,9 +453,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter():
# Mock the global handler instance in token_counter module
mock_handler = MagicMock()
mock_handler.handle_count_tokens_request = AsyncMock(
return_value={"input_tokens": 42}
)
mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42})
with patch(
"litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler",
@ -527,9 +512,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic():
# Mock the global handler instance in token_counter module
mock_handler = MagicMock()
mock_handler.handle_count_tokens_request = AsyncMock(
return_value={"input_tokens": 42}
)
mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42})
with patch(
"litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler",
@ -544,9 +527,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic():
mock_router.model_list = [
{
"model_name": "claude-3-5-sonnet",
"litellm_params": {
"model": "anthropic/claude-3-5-sonnet-20241022"
},
"litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"},
"model_info": {},
}
]
@ -555,9 +536,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic():
mock_router.async_get_available_deployment = AsyncMock(
return_value={
"model_name": "claude-3-5-sonnet",
"litellm_params": {
"model": "anthropic/claude-3-5-sonnet-20241022"
},
"litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"},
"model_info": {},
}
)
@ -592,9 +571,7 @@ async def test_factory_registration():
assert counter is not None
# Create test deployments
anthropic_deployment = {
"litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}
}
anthropic_deployment = {"litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}}
non_anthropic_deployment = {"litellm_params": {"model": "openai/gpt-4"}}
@ -630,9 +607,7 @@ async def test_bedrock_count_tokens_endpoint():
model_list=[
{
"model_name": "claude-bedrock",
"litellm_params": {
"model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0"
},
"litellm_params": {"model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0"},
}
]
)
@ -646,9 +621,7 @@ async def test_bedrock_count_tokens_endpoint():
}
# Test the mock handler directly to verify correct parameter extraction
await mock_count_tokens_handler(
request_data, {}, "anthropic.claude-3-sonnet-20240229-v1:0"
)
await mock_count_tokens_handler(request_data, {}, "anthropic.claude-3-sonnet-20240229-v1:0")
@pytest.mark.asyncio
@ -710,10 +683,7 @@ async def test_vertex_ai_anthropic_token_counting():
assert call_args is not None
assert call_args.kwargs["model"] == "claude-3-5-sonnet-20241022"
assert "messages" in call_args.kwargs["request_data"]
assert (
call_args.kwargs["request_data"]["messages"][0]["content"]
== "Hello Claude on Vertex AI! How are you?"
)
assert call_args.kwargs["request_data"]["messages"][0]["content"] == "Hello Claude on Vertex AI! How are you?"
# Validate response structure
assert response.model_used == "claude-3-5-sonnet-20241022"
@ -746,9 +716,7 @@ def test_vertex_ai_partner_models_token_counting_endpoint(vertex_location):
if vertex_location == "global":
assert endpoint.startswith("https://aiplatform.googleapis.com")
else:
assert endpoint.startswith(
f"https://{vertex_location}-aiplatform.googleapis.com"
)
assert endpoint.startswith(f"https://{vertex_location}-aiplatform.googleapis.com")
class _FailingRuntimeHandler(BedrockCountTokensHandler):
@ -834,20 +802,14 @@ async def test_bedrock_handler_httpx_error_status_code_propagation():
with patch.object(handler, "validate_count_tokens_request"):
with patch.object(handler, "_get_aws_region_name", return_value="us-west-2"):
with patch.object(
handler, "transform_anthropic_to_bedrock_count_tokens", return_value={}
):
with patch.object(handler, "transform_anthropic_to_bedrock_count_tokens", return_value={}):
with patch.object(
handler,
"get_bedrock_count_tokens_endpoint",
return_value="https://example.com",
):
with patch.object(
handler, "_sign_request", return_value=({}, "{}")
):
with patch(
"litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client"
) as mock_client:
with patch.object(handler, "_sign_request", return_value=({}, "{}")):
with patch("litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client") as mock_client:
mock_async_client = AsyncMock()
mock_async_client.post = AsyncMock(side_effect=http_error)
mock_client.return_value = mock_async_client
@ -856,9 +818,7 @@ async def test_bedrock_handler_httpx_error_status_code_propagation():
await handler.handle_count_tokens_request(
request_data={
"model": "test",
"messages": [
{"role": "user", "content": "hello"}
],
"messages": [{"role": "user", "content": "hello"}],
},
litellm_params={},
resolved_model="anthropic.claude-3-sonnet",
@ -866,10 +826,7 @@ async def test_bedrock_handler_httpx_error_status_code_propagation():
assert exc_info.value.status_code == 403
# Message should be the raw response text
assert (
exc_info.value.message
== "Forbidden - Invalid credentials"
)
assert exc_info.value.message == "Forbidden - Invalid credentials"
@pytest.mark.asyncio
@ -897,9 +854,7 @@ async def test_token_counter_httpx_status_error_raises_proxy_exception():
mock_counter.count_tokens = AsyncMock(side_effect=http_error)
# Save originals
original_get_provider_token_counter = (
litellm.proxy.proxy_server._get_provider_token_counter
)
original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter
original_router = litellm.proxy.proxy_server.llm_router
try:
@ -907,9 +862,7 @@ async def test_token_counter_httpx_status_error_raises_proxy_exception():
def mock_get_provider_token_counter(deployment, model_to_use):
return (mock_counter, "claude-4-6-sonnet", "vertex_ai")
litellm.proxy.proxy_server._get_provider_token_counter = (
mock_get_provider_token_counter
)
litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter
mock_router = MagicMock()
mock_router.async_get_available_deployment = AsyncMock(
@ -937,9 +890,7 @@ async def test_token_counter_httpx_status_error_raises_proxy_exception():
assert exc_info.value.type == "token_counting_error"
assert exc_info.value.param == "model"
finally:
litellm.proxy.proxy_server._get_provider_token_counter = (
original_get_provider_token_counter
)
litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter
litellm.proxy.proxy_server.llm_router = original_router
@ -975,9 +926,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled():
# Save original value and function
original_disable = litellm.disable_token_counter
original_get_provider_token_counter = (
litellm.proxy.proxy_server._get_provider_token_counter
)
original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter
try:
litellm.disable_token_counter = True
@ -991,9 +940,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled():
def mock_get_provider_token_counter(deployment, model_to_use):
return (mock_counter, "anthropic.claude-3-sonnet", "bedrock")
litellm.proxy.proxy_server._get_provider_token_counter = (
mock_get_provider_token_counter
)
litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter
with pytest.raises(ProxyException) as exc_info:
await token_counter(
@ -1008,9 +955,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled():
assert "Rate limit exceeded" in exc_info.value.message
finally:
litellm.disable_token_counter = original_disable
litellm.proxy.proxy_server._get_provider_token_counter = (
original_get_provider_token_counter
)
litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter
@pytest.mark.asyncio
@ -1045,9 +990,7 @@ async def test_proxy_token_counter_error_raises_when_enabled():
# Save original value and function
original_disable = litellm.disable_token_counter
original_get_provider_token_counter = (
litellm.proxy.proxy_server._get_provider_token_counter
)
original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter
try:
litellm.disable_token_counter = False
@ -1061,9 +1004,7 @@ async def test_proxy_token_counter_error_raises_when_enabled():
def mock_get_provider_token_counter(deployment, model_to_use):
return (mock_counter, "anthropic.claude-3-sonnet", "bedrock")
litellm.proxy.proxy_server._get_provider_token_counter = (
mock_get_provider_token_counter
)
litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter
with patch("litellm.token_counter") as mock_local_counter:
with pytest.raises(ProxyException) as exc_info:
@ -1081,9 +1022,7 @@ async def test_proxy_token_counter_error_raises_when_enabled():
mock_local_counter.assert_not_called()
finally:
litellm.disable_token_counter = original_disable
litellm.proxy.proxy_server._get_provider_token_counter = (
original_get_provider_token_counter
)
litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter
@pytest.mark.asyncio
@ -1109,10 +1048,10 @@ async def test_anthropic_endpoint_returns_anthropic_error_format():
original_read_request_body = anthropic_endpoints.read_request_body
anthropic_endpoints.read_request_body = mock_read_request_body
original_token_counter = proxy_server.token_counter
original_token_counter = proxy_server.count_request_tokens
# Mock token_counter to raise ProxyException with Bedrock-style error
async def mock_token_counter_error(request, call_endpoint=False):
async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None):
raise ProxyException(
message='{"detail":{"message":"Input is too long for requested model."}}',
type="token_counting_error",
@ -1120,7 +1059,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format():
code=400,
)
proxy_server.token_counter = mock_token_counter_error
proxy_server.count_request_tokens = mock_token_counter_error
try:
with pytest.raises(HTTPException) as exc_info:
@ -1136,7 +1075,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format():
assert detail["error"]["message"] == "Input is too long for requested model."
finally:
anthropic_endpoints.read_request_body = original_read_request_body
proxy_server.token_counter = original_token_counter
proxy_server.count_request_tokens = original_token_counter
@pytest.mark.asyncio
@ -1161,10 +1100,10 @@ async def test_anthropic_endpoint_403_permission_error_format():
original_read_request_body = anthropic_endpoints.read_request_body
anthropic_endpoints.read_request_body = mock_read_request_body
original_token_counter = proxy_server.token_counter
original_token_counter = proxy_server.count_request_tokens
# Mock token_counter to raise ProxyException with 403 error
async def mock_token_counter_error(request, call_endpoint=False):
async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None):
raise ProxyException(
message='{"Message":"Bearer Token has expired"}',
type="token_counting_error",
@ -1172,7 +1111,7 @@ async def test_anthropic_endpoint_403_permission_error_format():
code=403,
)
proxy_server.token_counter = mock_token_counter_error
proxy_server.count_request_tokens = mock_token_counter_error
try:
with pytest.raises(HTTPException) as exc_info:
@ -1186,7 +1125,7 @@ async def test_anthropic_endpoint_403_permission_error_format():
assert detail["error"]["message"] == "Bearer Token has expired"
finally:
anthropic_endpoints.read_request_body = original_read_request_body
proxy_server.token_counter = original_token_counter
proxy_server.count_request_tokens = original_token_counter
@pytest.mark.asyncio
@ -1211,10 +1150,10 @@ async def test_anthropic_endpoint_429_rate_limit_error_format():
original_read_request_body = anthropic_endpoints.read_request_body
anthropic_endpoints.read_request_body = mock_read_request_body
original_token_counter = proxy_server.token_counter
original_token_counter = proxy_server.count_request_tokens
# Mock token_counter to raise ProxyException with 429 error
async def mock_token_counter_error(request, call_endpoint=False):
async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None):
raise ProxyException(
message="Rate limit exceeded",
type="token_counting_error",
@ -1222,7 +1161,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format():
code=429,
)
proxy_server.token_counter = mock_token_counter_error
proxy_server.count_request_tokens = mock_token_counter_error
try:
with pytest.raises(HTTPException) as exc_info:
@ -1236,4 +1175,4 @@ async def test_anthropic_endpoint_429_rate_limit_error_format():
assert detail["error"]["message"] == "Rate limit exceeded"
finally:
anthropic_endpoints.read_request_body = original_read_request_body
proxy_server.token_counter = original_token_counter
proxy_server.count_request_tokens = original_token_counter

2
tests/unit/test_count_tokens_public_api.py
View file

@ -97,7 +97,7 @@ def test_acount_tokens_with_tools():
assert result.total_tokens == 30
mock_handler.assert_called_once()
call_kwargs = mock_handler.call_args
assert call_kwargs.kwargs.get("tools") == tools
assert list(call_kwargs.kwargs.get("tools") or ()) == tools
def test_acount_tokens_with_system():