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

fix(proxy): count the context_management edits /v1/messages applies on every bridged provider (#49234)

Browse source 
* fix(proxy): count the context_management edits /v1/messages applies on every bridged provider

* fix: apply context_management edits in count_tokens on Responses bridges, compaction blocks and Bedrock Claude

The Responses bridge now applies the same edits inference applies for every
provider, the chat bridge cuts the history at a client compaction block, and
Bedrock Claude forwards the edits AWS accepts with their beta to the mantle
and runtime count APIs instead of dropping them.

* refactor(count_tokens): return count-view failures as values and raise them once in the caller

* fix(bedrock): keep the runtime's 400 when mantle does not serve the model it retries

* fix(count_tokens): forward context edits on Vertex and Azure Claude counts and keep Bedrock's mantle answer without edits

* fix(proxy): count edited bodies on native routes without a count API, keep counts on editor failures

* fix(proxy): keep the Bedrock pass-through count shape and provider 400s on empty bridged bodies

The pass-through count picks Converse or InvokeModel from the content again, and the edited InvokeModel count wraps string turns in text blocks. A bridged view with no messages sends the body as before. Azure and Vertex load the beta config off the event loop before a count with edits, and a model with no deployment reads drop_params from the request and router defaults.

* fix(bedrock): send an empty-string turn as is when a count forwards context edits

Wrapping an empty string into a text block made AWS reject the count with 'text content blocks must be non-empty', while base and /v1/messages accept the empty string. Only non-empty string turns are wrapped now.

* test(azure_ai): compare the count's betas with the betas Azure inference sends

* test(azure_ai): record the Azure inference request with respx instead of appending to a list

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
tessellark 2026-11-02 10:38:35 -08:00 • committed by GitHub
parent 97cd5d0c6a
commit a2ab27373f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
22 changed files with 1691 additions and 393 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

138
litellm/llms/anthropic/messages/context_management/inference_view.py Normal file
View file

@ -0,0 +1,138 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Final
from pydantic import JsonValue, TypeAdapter, ValidationError
from typing_extensions import assert_never
from litellm.types.llms.anthropic import AnthropicMessagesTextParam
from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat
@dataclass(frozen=True, slots=True)
class InferenceRequestView:
messages: Sequence[Mapping[str, object]]
system: object | None
@dataclass(frozen=True, slots=True)
class InferenceRequestViewError:
error_message: str
status_code: int
@dataclass(frozen=True, slots=True)
class CountContextEdits:
spec: JsonValue
model: str
tools: tuple[Mapping[str, object], ...] | None
drop_invalid_edits: bool = False
_CONTEXT_MANAGEMENT_SPEC: Final[TypeAdapter[Mapping[str, object] | tuple[Mapping[str, object], ...]]] = TypeAdapter(
Mapping[str, object] | tuple[Mapping[str, object], ...]
)
_SYSTEM_BLOCKS: Final = TypeAdapter(tuple[AnthropicMessagesTextParam, ...])
def count_context_edits(
context_management: CountContextManagement | None,
tools: Sequence[Mapping[str, object]] | None,
litellm_params: Mapping[str, object],
fallback_model: str,
) -> CountContextEdits:
model: Final = litellm_params.get("model")
return CountContextEdits(
spec=None if context_management is None else context_management.spec,
model=model if isinstance(model, str) else fallback_model,
tools=None if tools is None else tuple(tools),
drop_invalid_edits=context_management is not None and context_management.drop_invalid_edits,
)
def _typed_spec(spec: JsonValue) -> Mapping[str, object] | tuple[Mapping[str, object], ...] | None:
try:
return _CONTEXT_MANAGEMENT_SPEC.validate_python(spec)
except ValidationError:
return None
def _typed_system(system: object | None) -> str | tuple[AnthropicMessagesTextParam, ...] | None | ValidationError:
if system is None or isinstance(system, str):
return system
try:
return _SYSTEM_BLOCKS.validate_python(system)
except ValidationError as error:
return error
def _view_without_edits(
messages: Sequence[Mapping[str, object]],
system: object | None,
typed_system: str | Sequence[Mapping[str, object]] | None,
) -> InferenceRequestView:
from litellm.llms.anthropic.messages.context_management.editors.compact import (
apply_client_compaction_block_history,
)
history: Final = apply_client_compaction_block_history(
messages=[dict(message) for message in messages],
system=typed_system
if typed_system is None or isinstance(typed_system, str)
else [dict(block) for block in typed_system],
)
if history is None:
return InferenceRequestView(messages=messages, system=system)
return InferenceRequestView(messages=history.messages, system=history.system)
async def _anthropic_inference_view(
messages: Sequence[Mapping[str, object]], system: object | None, edits: CountContextEdits
) -> InferenceRequestView | InferenceRequestViewError:
from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request
from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError
from litellm.llms.anthropic.messages.native.utils import normalize_native_messages
typed_system: Final = _typed_system(system)
if isinstance(typed_system, ValidationError):
return InferenceRequestView(messages=messages, system=system)
normalized: Final = normalize_native_messages(messages)
spec: Final = _typed_spec(edits.spec)
try:
prepared: Final = await prepare_context_managed_request(
model=edits.model,
messages=normalized,
tools=None if edits.tools is None else [dict(tool) for tool in edits.tools],
system=[dict(block) for block in typed_system] if isinstance(typed_system, tuple) else typed_system,
context_management_spec=(
None if spec is None else dict(spec) if isinstance(spec, Mapping) else [dict(edit) for edit in spec]
),
litellm_metadata=None,
additional_drop_params=None,
llm_router=None,
count_only=True,
drop_invalid_edits=edits.drop_invalid_edits,
)
except AnthropicContextManagementError as error:
if error.status_code >= 500:
return _view_without_edits(normalized, system, typed_system)
return InferenceRequestViewError(error_message=error.message, status_code=error.status_code)
if prepared is None:
return InferenceRequestView(messages=normalized, system=system)
return InferenceRequestView(messages=prepared.messages, system=prepared.system)
async def inference_request_view(
messages: Sequence[Mapping[str, object]],
system: object | None,
messages_format: TokenCountMessagesFormat,
edits: CountContextEdits,
) -> InferenceRequestView | InferenceRequestViewError:
"""Apply the history rewrites and context edits /v1/messages runs before it converts a body for another provider."""
match messages_format:
case "openai":
return InferenceRequestView(messages=messages, system=system)
case "anthropic":
return await _anthropic_inference_view(messages, system, edits)
case _:
assert_never(messages_format)

6
litellm/llms/azure_ai/anthropic/count_tokens/handler.py
View file

@ -12,6 +12,7 @@ from pydantic import JsonValue
import litellm
from litellm._logging import verbose_logger
from litellm.anthropic_beta_headers_manager import aload_beta_headers_config
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.llms.anthropic.common_utils import AnthropicError
@ -39,6 +40,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
tools: list[dict[str, JsonValue]] | None = None,
system: object = None,
optional_params: Mapping[str, JsonValue] | None = None,
context_management: JsonValue = None,
) -> dict[str, object]:
"""
Handle a CountTokens request using httpx with Azure authentication.
@ -70,6 +72,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
tools=tools,
system=system,
optional_params=optional_params,
context_management=context_management,
drop_params=drops_unsupported_params((litellm_params or {}).get("drop_params")),
)
@ -80,11 +83,14 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
verbose_logger.debug("Making request to: %s", endpoint_url)
if context_management is not None:
await aload_beta_headers_config()
# Get required headers with Azure authentication
headers: Final = self.get_required_headers(
api_key=api_key,
litellm_params=litellm_params,
messages=messages,
context_management=context_management,
)
if not {"x-api-key", "Authorization"} & headers.keys():
raise AnthropicError(

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

@ -135,6 +135,7 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
tools=tools,
system=system,
optional_params=count_options,
context_management=None if context_management is None else context_management.spec,
)
if result is not None:

17
litellm/llms/azure_ai/anthropic/count_tokens/transformation.py
View file

@ -8,9 +8,16 @@ from collections.abc import Sequence
from typing import Final
from urllib.parse import urlsplit
from pydantic import JsonValue
from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers
from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.anthropic.common_utils import count_tokens_feature_betas, merge_anthropic_beta_headers
from litellm.llms.anthropic.common_utils import (
context_management_betas,
count_tokens_feature_betas,
merge_anthropic_beta_headers,
)
from litellm.llms.anthropic.count_tokens.transformation import (
AnthropicCountTokensConfig,
)
@ -31,6 +38,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
api_key: str | None,
litellm_params: dict[str, object] | None = None,
messages: Sequence[object] = (),
context_management: JsonValue = None,
) -> dict[str, str]:
"""
Get the required headers for the Azure AI Anthropic CountTokens API.
@ -42,6 +50,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
api_key: The Azure AI API key
litellm_params: Optional LiteLLM parameters for additional auth config
messages: The request messages, read for betas the counted features need
context_management: The edit spec the count forwards, read for the betas its edits need
Returns:
Dictionary of required headers with both x-api-key and Azure authentication
@ -51,7 +60,11 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
"Content-Type": "application/json",
"anthropic-version": "2023-06-01",
"anthropic-beta": merge_anthropic_beta_headers(
ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, count_tokens_feature_betas("azure_ai", {}, messages)
ANTHROPIC_TOKEN_COUNTING_BETA_VERSION,
(
*count_tokens_feature_betas("azure_ai", {}, messages),
*filter_and_transform_beta_headers(list(context_management_betas(context_management)), "azure_ai"),
),
),
**({"x-api-key": api_key} if api_key else {}),
}

7
litellm/llms/bedrock/common_utils.py
View file

@ -43,6 +43,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned
from litellm.secret_managers.main import get_secret, get_secret_str
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams
if TYPE_CHECKING:
@ -50,6 +51,12 @@ if TYPE_CHECKING:
from litellm.types.llms.openai import AllMessageValues
BEDROCK_INVOKE_TRACE_FIELD: Final = "amazon-bedrock-trace"
BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS: Final[Mapping[str, str]] = MappingProxyType(
{
"compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
"clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
}
)
_ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs"

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

@ -13,9 +13,14 @@ from typing_extensions import assert_never
import litellm
from litellm._logging import verbose_logger
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.common_utils import context_management_betas, context_management_with_edits
from litellm.llms.anthropic.count_tokens.transformation import ChatCountTranslationError, anthropic_count_fields
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model
from litellm.llms.bedrock.common_utils import (
BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS,
BedrockError,
get_bedrock_base_model,
)
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
@ -74,6 +79,16 @@ def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model:
)
def _mantle_retry_outcome(
runtime: CountTokensOutcome, mantle: CountTokensOutcome, carried_edits: bool
) -> CountTokensOutcome:
match runtime, mantle:
case CountTokensFailure(status_code=400), CountTokensFailure(status_code=404) if carried_edits:
return runtime
case _:
return mantle
def _request_data(
model: str, messages: Sequence[Mapping[str, object]], tools: Sequence[Mapping[str, object]] | None, system: object
) -> Mapping[str, object]:
@ -85,6 +100,33 @@ def _request_data(
}
def _invoke_model_message(message: Mapping[str, object]) -> Mapping[str, object]:
content: Final = message.get("content")
return (
{**message, "content": [{"type": "text", "text": content}]} if isinstance(content, str) and content else message
)
def _invoke_supports_edit(edit: JsonValue) -> bool:
edit_type: Final = edit.get("type") if isinstance(edit, dict) else None
return isinstance(edit_type, str) and edit_type in BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS
def _invoke_context_management(context_management: CountContextManagement | None) -> Mapping[str, JsonValue]:
spec: Final = context_management_with_edits(None if context_management is None else context_management.spec)
edits: Final = spec.get("edits") if isinstance(spec, dict) else None
supported_edits: Final = tuple(filter(_invoke_supports_edit, edits if isinstance(edits, list) else ()))
if not isinstance(spec, dict) or not supported_edits:
return MappingProxyType({})
supported_spec: Final[JsonValue] = {**spec, "edits": list(supported_edits)}
return MappingProxyType(
{
"context_management": supported_spec,
"anthropic_beta": list(context_management_betas(supported_spec)),
}
)
class BedrockTokenCounter(BaseTokenCounter):
"""Counts on the configured Bedrock endpoint, with Mantle retry when Runtime rejects a Claude model"""
@ -167,11 +209,19 @@ class BedrockTokenCounter(BaseTokenCounter):
)
if isinstance(translated, ChatCountTranslationError):
return translated.response(request_model, model_to_use, primary_tokenizer)
runtime_request: Final = (
_request_data(model_to_use, messages, tools, system)
if translated is None
else _request_data(model_to_use, translated.messages, translated.tools, translated.system)
native_context_management: Final = _invoke_context_management(
context_management if is_mantle_claude_model(resolved_model) else None
)
counted_messages: Final = messages if translated is None else translated.messages
runtime_request: Final = {
**_request_data(
model_to_use,
tuple(map(_invoke_model_message, counted_messages)) if native_context_management else counted_messages,
tools if translated is None else translated.tools,
system if translated is None else translated.system,
),
**native_context_management,
}
mantle_request: Final = runtime_request if translated is None else {**runtime_request, **translated.options}
primary: Final = (
@ -184,8 +234,12 @@ class BedrockTokenCounter(BaseTokenCounter):
)
)
outcome: Final = (
await self._count_with(
self._mantle_handler, MANTLE_TOKENIZER_TYPE, mantle_request, litellm_params, resolved_model
_mantle_retry_outcome(
primary,
await self._count_with(
self._mantle_handler, MANTLE_TOKENIZER_TYPE, mantle_request, litellm_params, resolved_model
),
carried_edits=bool(native_context_management),
)
if not use_mantle and _runtime_rejected_claude_model(primary, resolved_model)
else primary

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

@ -8,7 +8,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.anthropic.common_utils import tools_without_custom_field
from litellm.llms.anthropic.common_utils import context_management_betas, tools_without_custom_field
from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, AnthropicCountTokensConfig
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url
@ -22,6 +22,7 @@ class MantleCountTokensRequest(TypedDict):
messages: ReadOnly[list[dict[str, JsonValue]]]
system: ReadOnly[NotRequired[JsonValue]]
tools: ReadOnly[NotRequired[list[dict[str, JsonValue]]]]
context_management: ReadOnly[NotRequired[JsonValue]]
_COUNT_REQUEST: Final = TypeAdapter(MantleCountTokensRequest)
@ -74,13 +75,19 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM):
optional_params=_COUNT_OPTIONS.validate_python(
{name: value for name, value in request_data.items() if name in COUNT_TOKEN_OPTION_NAMES}
),
context_management=request.get("context_management"),
)
betas: Final = context_management_betas(request.get("context_management"))
verbose_logger.debug("Making bedrock-mantle count_tokens request to: %s", endpoint_url)
api_key: Final = litellm_params.get("api_key")
signed_headers, signed_body = await run_aws_signing(
self._sign_request,
service_name="bedrock",
headers={"Content-Type": "application/json", "anthropic-version": MANTLE_ANTHROPIC_VERSION},
headers={
"Content-Type": "application/json",
"anthropic-version": MANTLE_ANTHROPIC_VERSION,
**({"anthropic-beta": ",".join(betas)} if betas else {}),
},
optional_params=litellm_params,
request_data=body,
api_base=endpoint_url,

8
litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py
View file

@ -37,6 +37,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation
bedrock_invoke_guardrail_headers,
)
from litellm.llms.bedrock.common_utils import (
BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS,
BedrockError,
apply_bedrock_invoke_structured_output,
bedrock_supports_tool_search,
@ -499,12 +500,7 @@ class AmazonAnthropicClaudeMessagesConfig(
# Bedrock InvokeModel DOES support ``clear_tool_uses_20250919`` under the
# ``context-management-2025-06-27`` beta. AWS docs:
# https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-tool-use.md
_BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = MappingProxyType(
{
"compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
"clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
}
)
_BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS
@classmethod
def _filter_context_management_for_bedrock_invoke(

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

@ -537,13 +537,16 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter):
) -> TokenCountResponse | None:
import copy
from litellm.llms.anthropic.messages.context_management.inference_view import (
InferenceRequestViewError,
count_context_edits,
inference_request_view,
)
from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter
from litellm.llms.vertex_ai.count_tokens.context import (
VertexGeminiCountContextError,
count_context_edits,
count_context_error_response,
gemini_count_contents,
inference_request_view,
malformed_tool_result_text_count_error,
provider_error_response,
unconvertible_messages_error,
@ -556,7 +559,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter):
messages_format,
count_context_edits(context_management, tools, (deployment or {}).get("litellm_params", {}), model_to_use),
)
if isinstance(request_view, VertexGeminiCountContextError):
if isinstance(request_view, InferenceRequestViewError):
return count_context_error_response(request_view, request_model, model_to_use)
count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools, model=model_to_use)
if isinstance(count_context, VertexGeminiCountContextError):

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

@ -1458,6 +1458,7 @@ class VertexAITokenCounter(BaseTokenCounter):
) -> TokenCountResponse | None:
import copy
from litellm.llms.anthropic.common_utils import context_management_with_edits
from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_publisher
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
VertexAIError as PartnerVertexAIError,
@ -1539,6 +1540,9 @@ class VertexAITokenCounter(BaseTokenCounter):
system=system,
tools=tools,
count_options=count_options,
context_management=(
None if context_management is None else context_management_with_edits(context_management.spec)
),
)
except (PartnerVertexAIError, httpx.HTTPStatusError) as e:
status_code: Final = e.response.status_code
@ -1565,12 +1569,15 @@ class VertexAITokenCounter(BaseTokenCounter):
original_response=result,
)
else:
from litellm.llms.anthropic.messages.context_management.inference_view import (
InferenceRequestViewError,
count_context_edits,
inference_request_view,
)
from litellm.llms.vertex_ai.count_tokens.context import (
VertexGeminiCountContextError,
count_context_edits,
count_context_error_response,
gemini_count_contents,
inference_request_view,
malformed_tool_result_text_count_error,
provider_error_response,
unconvertible_messages_error,
@ -1584,7 +1591,7 @@ class VertexAITokenCounter(BaseTokenCounter):
messages_format,
count_context_edits(context_management, tools, count_tokens_params_request, model_to_use),
)
if isinstance(request_view, VertexGeminiCountContextError):
if isinstance(request_view, InferenceRequestViewError):
return count_context_error_response(request_view, request_model, model_to_use)
count_context: Final = vertex_gemini_count_context(
system=request_view.system, tools=tools, model=model_to_use

129
litellm/llms/vertex_ai/count_tokens/context.py
View file

@ -11,10 +11,11 @@ from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_
from litellm.types.llms.anthropic import AnthropicMessagesTextParam
from litellm.types.llms.openai import AllMessageValues
from litellm.types.llms.vertex_ai import ContentType
from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse
from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
if TYPE_CHECKING:
from litellm.exceptions import APIError
from litellm.llms.anthropic.messages.context_management.inference_view import InferenceRequestViewError
@dataclass(frozen=True, slots=True)
@ -152,7 +153,7 @@ def malformed_tool_result_text_count_error(
def count_context_error_response(
error: VertexGeminiCountContextError, request_model: str, model_used: str
error: "VertexGeminiCountContextError | InferenceRequestViewError", request_model: str, model_used: str
) -> TokenCountResponse:
return TokenCountResponse(
total_tokens=0,
@ -178,130 +179,6 @@ def provider_error_response(error: "APIError", request_model: str, model_used: s
)
@dataclass(frozen=True, slots=True)
class InferenceRequestView:
messages: Sequence[Mapping[str, object]]
system: object | None
@dataclass(frozen=True, slots=True)
class CountContextEdits:
spec: JsonValue
model: str
tools: tuple[Mapping[str, object], ...] | None
drop_invalid_edits: bool = False
_CONTEXT_MANAGEMENT_SPEC: Final[TypeAdapter[Mapping[str, object] | tuple[Mapping[str, object], ...]]] = TypeAdapter(
Mapping[str, object] | tuple[Mapping[str, object], ...]
)
def count_context_edits(
context_management: CountContextManagement | None,
tools: Sequence[Mapping[str, object]] | None,
litellm_params: Mapping[str, object],
fallback_model: str,
) -> CountContextEdits:
model: Final = litellm_params.get("model")
return CountContextEdits(
spec=None if context_management is None else context_management.spec,
model=model if isinstance(model, str) else fallback_model,
tools=None if tools is None else tuple(tools),
drop_invalid_edits=context_management is not None and context_management.drop_invalid_edits,
)
def _typed_spec(spec: JsonValue) -> Mapping[str, object] | tuple[Mapping[str, object], ...] | None:
try:
return _CONTEXT_MANAGEMENT_SPEC.validate_python(spec)
except ValidationError:
return None
def _typed_system(
system: object | None,
) -> str | tuple[AnthropicMessagesTextParam, ...] | None | VertexGeminiCountContextError:
if system is None or isinstance(system, str):
return system
try:
return _SYSTEM_BLOCKS.validate_python(system)
except ValidationError:
return _INVALID_SYSTEM
def _view_without_edits(
messages: Sequence[Mapping[str, object]],
system: object | None,
typed_system: str | Sequence[Mapping[str, object]] | None,
) -> InferenceRequestView:
from litellm.llms.anthropic.messages.context_management.editors.compact import (
apply_client_compaction_block_history,
)
history: Final = apply_client_compaction_block_history(
messages=[dict(message) for message in messages],
system=typed_system
if typed_system is None or isinstance(typed_system, str)
else [dict(block) for block in typed_system],
)
if history is None:
return InferenceRequestView(messages=messages, system=system)
return InferenceRequestView(messages=history.messages, system=history.system)
async def _anthropic_inference_view(
messages: Sequence[Mapping[str, object]], system: object | None, edits: CountContextEdits
) -> InferenceRequestView | VertexGeminiCountContextError:
from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request
from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError
from litellm.llms.anthropic.messages.native.utils import normalize_native_messages
typed_system: Final = _typed_system(system)
if isinstance(typed_system, VertexGeminiCountContextError):
return typed_system
normalized: Final = normalize_native_messages(messages)
spec: Final = _typed_spec(edits.spec)
try:
prepared: Final = await prepare_context_managed_request(
model=edits.model,
messages=normalized,
tools=None if edits.tools is None else [dict(tool) for tool in edits.tools],
system=[dict(block) for block in typed_system] if isinstance(typed_system, tuple) else typed_system,
context_management_spec=(
None if spec is None else dict(spec) if isinstance(spec, Mapping) else [dict(edit) for edit in spec]
),
litellm_metadata=None,
additional_drop_params=None,
llm_router=None,
count_only=True,
drop_invalid_edits=edits.drop_invalid_edits,
)
except AnthropicContextManagementError as error:
if error.status_code >= 500:
return _view_without_edits(normalized, system, typed_system)
return VertexGeminiCountContextError(error_message=error.message, status_code=error.status_code)
if prepared is None:
return InferenceRequestView(messages=normalized, system=system)
return InferenceRequestView(messages=prepared.messages, system=prepared.system)
async def inference_request_view(
messages: Sequence[Mapping[str, object]],
system: object | None,
messages_format: TokenCountMessagesFormat,
edits: CountContextEdits,
) -> InferenceRequestView | VertexGeminiCountContextError:
"""Apply the history rewrites and context edits /v1/messages runs before it converts a body for Gemini."""
match messages_format:
case "openai":
return InferenceRequestView(messages=messages, system=system)
case "anthropic":
return await _anthropic_inference_view(messages, system, edits)
case _:
assert_never(messages_format)
def _chat_messages(
messages: Sequence[Mapping[str, object]],
messages_format: TokenCountMessagesFormat,

17
litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py
View file

@ -9,12 +9,14 @@ their respective publisher-specific count-tokens endpoints.
from collections.abc import Mapping, Sequence
from typing import Any, Final
from pydantic import TypeAdapter
from pydantic import JsonValue, TypeAdapter
from litellm.anthropic_beta_headers_manager import aload_beta_headers_config, filter_and_transform_beta_headers
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.anthropic.common_utils import (
OBJECT_LIST_ADAPTER,
context_management_betas,
count_tokens_feature_betas,
is_claude_model_target,
tools_without_custom_field,
@ -31,6 +33,7 @@ _COUNT_TOKENS_PUBLISHERS: Final = (
_REQUEST_MESSAGES: Final = TypeAdapter(Sequence[object])
_CONTEXT_MANAGEMENT: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
def count_tokens_request_body(
@ -55,8 +58,14 @@ def count_tokens_publisher(model: str) -> str | None:
def count_tokens_headers(access_token: str, request_data: Mapping[str, object]) -> Mapping[str, str]:
authorization: Final = {"Authorization": f"Bearer {access_token}"}
betas: Final = count_tokens_feature_betas(
"vertex_ai", request_data, _REQUEST_MESSAGES.validate_python(request_data["messages"])
betas: Final = (
*count_tokens_feature_betas(
"vertex_ai", request_data, _REQUEST_MESSAGES.validate_python(request_data["messages"])
),
*filter_and_transform_beta_headers(
list(context_management_betas(_CONTEXT_MANAGEMENT.validate_python(request_data.get("context_management")))),
"vertex_ai",
),
)
return {**authorization, "anthropic-beta": ",".join(betas)} if betas else authorization
@ -193,6 +202,8 @@ class VertexAIPartnerModelsTokenCounter(VertexBase):
api_base=litellm_params.get("api_base"),
)
if request_data.get("context_management") is not None:
await aload_beta_headers_config()
headers: Final = dict(count_tokens_headers(access_token, request_data))
# Get async HTTP client

11
litellm/llms/vertex_ai/vertex_ai_partner_models/main.py
View file

@ -268,6 +268,7 @@ class VertexAIPartnerModels(VertexBase):
system: object | None = None,
tools: list[dict[str, object]] | None = None,
count_options: Mapping[str, object] = MappingProxyType({}),
context_management: object = None,
):
"""
Count tokens for Vertex AI partner models (Anthropic Claude, Mistral, etc.)
@ -302,7 +303,15 @@ class VertexAIPartnerModels(VertexBase):
"model": model,
"messages": messages,
**MappingProxyType(
{key: value for key, value in (("system", system), ("tools", tools)) if value is not None}
{
key: value
for key, value in (
("system", system),
("tools", tools),
("context_management", context_management),
)
if value is not None
}
),
**count_options,
}

42
litellm/messages/dispatch.py
View file

@ -75,9 +75,7 @@ def _public_request(
return native_call(legacy, args, _native_kwargs(fields, kwargs))
def _resolved(request: NativeCall) -> tuple[str, str | None]:
model: Final = str(request.resolved["model"])
provider: Final = optional_str(request.resolved.get("custom_llm_provider"))
def _resolved_model(model: str, provider: str | None) -> tuple[str, str | None]:
try:
resolved_model, resolved_provider, _, _ = get_llm_provider(model, provider)
except BadRequestError:
@ -85,21 +83,43 @@ def _resolved(request: NativeCall) -> tuple[str, str | None]:
return resolved_model, resolved_provider
def _target(request: NativeCall, model: str, provider: str | None) -> Target:
requested_model: Final = str(request.resolved["model"])
model_info: Final = request.resolved.get("model_info")
api_base: Final = request.resolved.get("api_base")
native_config: Final = main._native_messages_config( # pyright: ignore[reportPrivateUsage] # the handler's own passthrough lookup
model, provider, model_info, api_base if isinstance(api_base, str) else None
)
def _resolved(request: NativeCall) -> tuple[str, str | None]:
return _resolved_model(str(request.resolved["model"]), optional_str(request.resolved.get("custom_llm_provider")))
def _route_target(
requested_model: str,
model: str,
provider: str | None,
model_info: object,
api_base: str | None,
compacting: bool,
) -> Target:
native_config: Final = main._native_messages_config(model, provider, model_info, api_base) # pyright: ignore[reportPrivateUsage] # the handler's own passthrough lookup
if native_config is not None:
return Target.MESSAGES
compacting: Final = request.resolved.get("compaction") is not None
if not compacting and main._should_route_to_responses_api(provider, requested_model, model, model_info): # pyright: ignore[reportPrivateUsage] # the handler's own translation choice
return Target.RESPONSES
return Target.CHAT_COMPLETIONS
def _target(request: NativeCall, model: str, provider: str | None) -> Target:
api_base: Final = request.resolved.get("api_base")
return _route_target(
str(request.resolved["model"]),
model,
provider,
request.resolved.get("model_info"),
api_base if isinstance(api_base, str) else None,
compacting=request.resolved.get("compaction") is not None,
)
def messages_target(model: str, custom_llm_provider: str | None, model_info: object, api_base: str | None) -> Target:
resolved_model, resolved_provider = _resolved_model(model, custom_llm_provider)
return _route_target(model, resolved_model, resolved_provider, model_info, api_base, compacting=False)
_FEATURES: Final[detectors.Detectors] = MappingProxyType(
{
**detectors.COMMON,

190
litellm/proxy/proxy_server.py
View file

@ -33,7 +33,7 @@ from collections.abc import (
from collections.abc import (
Set as AbstractSet,
)
from dataclasses import dataclass
from dataclasses import dataclass, replace
from datetime import datetime, timedelta, timezone
from itertools import chain
from types import MappingProxyType, TracebackType, UnionType
@ -199,6 +199,7 @@ from litellm.router_strategy.auto_router.auto_router_tuning_baseline import (
snapshot_tuning_baselines,
tuning_limit_violation,
)
from litellm.rust_bridge.catalog import Target
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import (
PRICING_OVERRIDES_KEY,
@ -389,6 +390,12 @@ from litellm.litellm_core_utils.sensitive_data_masker import (
mask_sensitive_keys,
)
from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact
from litellm.llms.anthropic.messages.context_management.inference_view import (
InferenceRequestView,
InferenceRequestViewError,
count_context_edits,
inference_request_view,
)
from litellm.llms.anthropic.messages.convertible_blocks import (
deployment_drops_params,
drop_non_string_tool_result_texts,
@ -399,6 +406,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.openai.live import parse_live_session_start
from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.messages.dispatch import messages_target
from litellm.proxy._lazy_features import attach_lazy_features, reserve_lazy_slot
from litellm.proxy._types import *
from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolver, resolve_agent_access_group_ceiling
@ -14927,7 +14935,7 @@ async def _try_provider_token_count(
deployment: Mapping[str, object] | None,
request_model: str,
tools: list | None = None,
system: str | None = None,
system: object | None = None,
messages_format: TokenCountMessagesFormat = "openai",
output_format: JsonValue = None,
context_management: CountContextManagement | None = None,
@ -14997,6 +15005,7 @@ async def _provider_token_count(
request_model: str,
settings: TokenCountSettings,
local_estimate_allowed: bool,
forwarded_edits: bool = False,
) -> Union["TokenCountResponse", ProviderCountFailure, None]:
"""Count with the provider, using the token_count_model in its region when one is set.
@ -15020,7 +15029,7 @@ async def _provider_token_count(
pinned_count_model: Final = None if count_model is None or count_model.region is None else count_model
region_hint: Final = (
None
if pinned_count_model is None
if pinned_count_model is None or (forwarded_edits and counted.status_code == 400)
else f"This model counts tokens with {pinned_count_model.model} in {pinned_count_model.region}. "
"If that region is not enabled for the AWS account, enable it or set token_count_model in model_info"
)
@ -15282,6 +15291,9 @@ class _TokenCountDeployment(TypedDict, total=False):
_TOKEN_COUNT_DEPLOYMENT: Final = TypeAdapter(_TokenCountDeployment)
_COUNT_TOOLS: Final[TypeAdapter[tuple[Mapping[str, object], ...] | None]] = TypeAdapter(
tuple[Mapping[str, object], ...] | None
)
_ROUTED_DEPLOYMENT: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
@ -15351,12 +15363,15 @@ def _count_context_management(
model, general_settings_view(), user_model if isinstance(user_model, str) else None
)
router_defaults: Final = _router_default_litellm_params(llm_router)
sent_params: Final = tuple(
{**deployment["litellm_params"], **router_defaults, **request_litellm_params}
deployment_params: Final = tuple(
deployment["litellm_params"]
for deployment in llm_router.registry().deployments(
inference_model if isinstance(inference_model, str) else model, source="routable"
)
)
sent_params: Final = tuple(
{**params, **router_defaults, **request_litellm_params} for params in deployment_params or ({},)
)
if any(context_management_explicitly_dropped(params.get("additional_drop_params")) for params in sent_params):
return None
return CountContextManagement(
@ -15366,6 +15381,90 @@ def _count_context_management(
)
_GATEWAY_EDIT_ROUTES: Final = frozenset({Target.CHAT_COMPLETIONS, Target.RESPONSES})
def _messages_route(deployment: _TokenCountDeployment, request_model: str) -> Target:
params: Final = deployment.get("litellm_params", {})
return messages_target(
_non_empty_str(params.get("model")) or request_model,
_non_empty_str(params.get("custom_llm_provider")),
deployment.get("model_info"),
_deployment_api_base(deployment),
)
def _rejected_view(status_code: int, message: str) -> ProxyException:
return ProviderCountFailure(
status_code=status_code, message=message, rejected_before_provider_call=True
).to_proxy_exception()
async def _bridged_request_view(
messages: Sequence[Mapping[str, object]],
system: object,
tools: Sequence[Mapping[str, object]] | None,
messages_format: TokenCountMessagesFormat,
edits: CountContextManagement | None,
deployment: _TokenCountDeployment,
fallback_model: str,
) -> InferenceRequestView | InferenceRequestViewError | None:
view: Final = await inference_request_view(
messages,
system,
messages_format,
count_context_edits(edits, tools, deployment.get("litellm_params", {}), fallback_model),
)
return None if isinstance(view, InferenceRequestView) and not view.messages else view
async def _responses_request_view(
messages: Sequence[Mapping[str, object]],
system: object,
tools: Sequence[Mapping[str, object]] | None,
edits: CountContextManagement,
deployment: _TokenCountDeployment,
fallback_model: str,
) -> InferenceRequestView | InferenceRequestViewError:
from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError
from litellm.llms.anthropic.messages.responses_adapters.handler import count_view_messages
try:
edited: Final = await count_view_messages(
model=_non_empty_str(deployment.get("litellm_params", {}).get("model")) or fallback_model,
messages=messages,
tools=tools,
system=system,
context_management=edits,
)
except AnthropicContextManagementError as error:
if error.status_code >= 500:
return InferenceRequestView(messages=messages, system=system)
return InferenceRequestViewError(error_message=error.message, status_code=error.status_code)
return InferenceRequestView(messages=edited, system=system)
async def _gateway_request_view(
route: Target,
messages: Sequence[Mapping[str, object]],
system: object,
tools: Sequence[Mapping[str, object]] | None,
messages_format: TokenCountMessagesFormat,
edits: CountContextManagement | None,
deployment: _TokenCountDeployment,
fallback_model: str,
) -> InferenceRequestView | InferenceRequestViewError | None:
match route:
case Target.CHAT_COMPLETIONS:
return await _bridged_request_view(
messages, system, tools, messages_format, edits, deployment, fallback_model
)
case Target.RESPONSES if edits is not None:
return await _responses_request_view(messages, system, tools, edits, deployment, fallback_model)
case _:
return None
def _responses_routing_kwargs(responses_request: Mapping[str, object] | None) -> Mapping[str, object]:
previous_response_id: Final = None if responses_request is None else responses_request.get("previous_response_id")
return {"previous_response_id": previous_response_id} if isinstance(previous_response_id, str) else {}
@ -15559,6 +15658,40 @@ async def _count_request_tokens(
)
prompt_only: Final = bool(prompt) and conversation_empty
request_messages: Final = _counted_request_messages(messages, messages_format, count_deployment)
local_count_error: Final = malformed_tool_result_count_error(request_messages, messages_format, count_deployment)
malformed_messages: Final = (
None
if local_count_error is None
else ProxyException(message=local_count_error, type="token_counting_error", param="messages", code=400)
)
context_edits: Final = _count_context_management(context_management, request.model, request_litellm_params)
route: Final = (
_messages_route(deployment_view, request.model)
if request_messages is not None and messages_format == "anthropic"
else None
)
if context_edits is not None and route in _GATEWAY_EDIT_ROUTES and malformed_messages is not None:
raise malformed_messages
edited_request: Final = (
None
if route is None or request_messages is None
else await _gateway_request_view(
route,
request_messages,
system,
_COUNT_TOOLS.validate_python(tools),
messages_format,
context_edits,
deployment_view,
model_to_use,
)
)
if isinstance(edited_request, InferenceRequestViewError):
raise _rejected_view(edited_request.status_code, edited_request.error_message)
inference_messages: Final = (
request_messages if edited_request is None else [dict(message) for message in edited_request.messages]
)
inference_system: Final = system if edited_request is None else edited_request.system
contents_estimable: Final = not contents or (bool(messages) and _estimable_contents(contents))
prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES
routed_deployment: Final = _ROUTED_DEPLOYMENT.validate_python(deployment)
@ -15598,30 +15731,30 @@ async def _count_request_tokens(
if chat_count is not None
else [{"role": "user", "content": prompt}]
if prompt_as_message
else request_messages
else inference_messages
),
contents=None if prompt_as_message else contents,
deployment=count_deployment,
request_model=request.model,
tools=tools if chat_count is None else [dict(tool) for tool in chat_count.tools or ()] or None,
system=system,
system=inference_system,
messages_format=messages_format,
output_format=output_format,
context_management=_count_context_management(context_management, request.model, request_litellm_params),
context_management=None if edited_request is not None else context_edits,
chat_params=MappingProxyType({}) if chat_count is None else chat_count.params,
),
model_to_use=model_to_use,
request_model=request.model,
settings=settings,
local_estimate_allowed=litellm.disable_token_counter is not True and contents_estimable,
forwarded_edits=context_edits is not None and edited_request is None,
)
)
if isinstance(ladder, TokenCountResponse):
return ladder
local_count_error: Final = malformed_tool_result_count_error(request_messages, messages_format, count_deployment)
if local_count_error is not None:
raise ProxyException(message=local_count_error, type="token_counting_error", param="messages", code=400)
if malformed_messages is not None:
raise malformed_messages
# Check if token counter is disabled before fallback
if litellm.disable_token_counter is True:
@ -15632,6 +15765,27 @@ async def _count_request_tokens(
code=503,
)
local_view: Final = (
None
if route is not Target.MESSAGES or context_edits is None or request_messages is None
else await _bridged_request_view(
request_messages,
system,
_COUNT_TOOLS.validate_python(tools),
messages_format,
replace(context_edits, drop_invalid_edits=False),
deployment_view,
model_to_use,
)
)
if isinstance(local_view, InferenceRequestViewError):
raise _rejected_view(local_view.status_code, local_view.error_message)
counted_view: Final = edited_request if local_view is None else local_view
local_messages: Final = (
request_messages if counted_view is None else [dict(message) for message in counted_view.messages]
)
local_system: Final = system if counted_view is None else counted_view.system
# Default LiteLLM token counting
custom_tokenizer: CustomHuggingfaceTokenizer | None = None
if model_info is not None:
@ -15639,11 +15793,11 @@ async def _count_request_tokens(
CustomHuggingfaceTokenizer | None,
model_info.get("custom_tokenizer", None),
)
system_message: Final = _system_message(system)
system_message: Final = _system_message(local_system)
prompt_turn: Final = prompt_only and not (system_message is None and not tools)
typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes
Sequence[AllMessageValues] | None,
[{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else request_messages),
[{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else local_messages),
)
counted_messages: Final = (
typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages)
@ -15654,7 +15808,15 @@ async def _count_request_tokens(
framing: Final = settings.framing if counted_messages is not None and settings.framing is not None else 0
kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None
exact: Final = (
await _kimi_k3_count(request, messages_format, model_to_use, kimi_k3_messages, counted_tools)
await _kimi_k3_count(
request
if counted_view is None
else request.model_copy(update={"messages": local_messages, "system": local_system}),
messages_format,
model_to_use,
kimi_k3_messages,
counted_tools,
)
if kimi_k3_messages is not None and custom_tokenizer is None
else None
)

239
tests/unit/llms/anthropic/messages/context_management/test_inference_view.py Normal file
View file

@ -0,0 +1,239 @@
import logging
from typing import Final
from unittest.mock import patch
import pytest
from pydantic import JsonValue, TypeAdapter
import litellm
from litellm.llms.anthropic.messages.context_management.inference_view import (
CountContextEdits,
InferenceRequestView,
InferenceRequestViewError,
inference_request_view,
)
_NO_EDITS: Final = CountContextEdits(spec=None, model="gemini/gemini-2.5-flash", tools=None)
@pytest.mark.parametrize(
"system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system")
)
async def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None:
view: Final = await inference_request_view(
(
{"role": "user", "content": "first question " * 50},
{"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]},
{"role": "user", "content": "latest question"},
),
system,
"anthropic",
_NO_EDITS,
)
assert isinstance(view, InferenceRequestView)
assert view.messages == [{"role": "user", "content": "latest question"}]
assert "Summary of the earlier turns." in str(view.system)
assert "Be brief." in str(view.system)
async def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None:
snippet: Final = "LiteLLM ships a weekly release candidate and a stable line."
view: Final = await inference_request_view(
(
{"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,
}
],
},
],
},
),
None,
"anthropic",
_NO_EDITS,
)
assert isinstance(view, InferenceRequestView)
assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"])
assert [block["type"] for block in assistant_blocks] == ["text"]
assert snippet in assistant_blocks[0]["text"]
async def test_openai_view_is_the_request_unchanged() -> None:
messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},)
view: Final = await inference_request_view(messages, "Be brief.", "openai", _NO_EDITS)
assert isinstance(view, InferenceRequestView)
assert view.messages == messages
assert view.system == "Be brief."
_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180
_READ_FILE_TOOL: Final[dict[str, object]] = {
"name": "read_file",
"description": "Read a file",
"input_schema": {"type": "object", "properties": {"path": {"type": "string"}}},
}
_CLEAR_EDIT: Final[JsonValue] = {
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": 500},
"keep": {"type": "tool_uses", "value": 1},
}
_CLEAR_OLD_TOOL_RESULTS: Final[JsonValue] = {"edits": [_CLEAR_EDIT]}
def _agent_turns(rounds: int, first_round: int = 0) -> tuple[dict[str, JsonValue], ...]:
return (
*(
turn
for index in range(first_round, first_round + rounds)
for turn in (
{"role": "user", "content": f"Read file {index}."},
{
"role": "assistant",
"content": [
{"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}"}}
],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": f"toolu_{index}", "content": _TOOL_OUTPUT}],
},
)
), # comprehension-ok: flattens each round into its three turns
{"role": "user", "content": "Summarize what you read."},
)
def _tool_outputs_kept(view: InferenceRequestView) -> tuple[bool, ...]:
return tuple(
_TOOL_OUTPUT in str(message["content"]) for message in view.messages if "tool_result" in str(message["content"])
)
def _edits(spec: JsonValue) -> CountContextEdits:
return CountContextEdits(spec=spec, model="gemini/gemini-2.5-flash", tools=(_READ_FILE_TOOL,))
async def test_anthropic_view_clears_the_tool_results_inference_clears() -> None:
view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS))
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (False, False, True)
async def test_anthropic_view_applies_clear_edits_without_calling_the_compaction_summary_model() -> None:
spec: Final[JsonValue] = {
"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50000}}, _CLEAR_EDIT]
}
turns: Final = (
{"role": "user", "content": "first question " * 30000},
{"role": "assistant", "content": "first answer"},
*_agent_turns(3),
)
with (
patch(
"litellm.llms.anthropic.messages.context_management.editors.compact._read_summary_model_setting",
return_value="gemini/gemini-2.5-flash",
),
patch.object(litellm, "acompletion", side_effect=AssertionError("count_tokens must not summarize")) as summary,
):
view: Final = await inference_request_view(turns, None, "anthropic", _edits(spec))
assert summary.call_count == 0
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (False, False, True)
async def test_anthropic_view_counts_the_unedited_body_without_an_error_log_when_an_edit_fails(
caplog: pytest.LogCaptureFixture,
) -> None:
with (
patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")),
caplog.at_level(logging.DEBUG, logger="LiteLLM"),
):
view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS))
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (True, True, True)
assert any("skipping edits" in record.getMessage() for record in caplog.records)
assert not any(record.levelno >= logging.ERROR for record in caplog.records)
async def test_anthropic_view_counts_the_turns_after_a_compaction_block_when_an_edit_fails() -> None:
turns: Final = (
*_agent_turns(3)[:-1],
{"role": "assistant", "content": [{"type": "compaction", "content": "Read files 0 to 2."}]},
*_agent_turns(2, first_round=3),
)
with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")):
view: Final = await inference_request_view(turns, None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS))
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (True, True)
assert "Read files 0 to 2." in str(view.system)
async def test_anthropic_view_clears_only_the_turns_after_a_compaction_block_when_clear_runs_before_compact() -> None:
clear_before_compact: Final[JsonValue] = {
"edits": [
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "tool_uses", "value": 3},
"keep": {"type": "tool_uses", "value": 1},
},
{"type": "compact_20260112"},
]
}
turns: Final = (
*_agent_turns(3)[:-1],
{"role": "assistant", "content": [{"type": "compaction", "content": "Read files 0 to 2."}]},
*_agent_turns(2, first_round=3),
)
with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")):
view: Final = await inference_request_view(turns, None, "anthropic", _edits(clear_before_compact))
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (True, True)
assert "Read files 0 to 2." in str(view.system)
async def test_anthropic_view_leaves_a_system_that_is_not_text_blocks_and_its_turns_to_the_provider_check() -> None:
image_system: Final[JsonValue] = [{"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}]
view: Final = await inference_request_view(
_agent_turns(3), image_system, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)
)
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (True, True, True)
assert view.system == image_system
async def test_anthropic_view_rejects_a_compact_trigger_inference_rejects() -> None:
low_trigger: Final[JsonValue] = {
"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}]
}
view: Final = await inference_request_view(_agent_turns(1), None, "anthropic", _edits(low_trigger))
assert isinstance(view, InferenceRequestViewError)
assert view.status_code == 400

94
tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py
View file

@ -1,5 +1,7 @@
import json
import operator
import threading
from collections.abc import Mapping
from typing import Final
import httpx
@ -8,8 +10,9 @@ import respx
from pydantic import ValidationError
import litellm
from litellm import anthropic_beta_headers_manager
from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAnthropicTokenCounter
from litellm.types.utils import TokenCountResponse
from litellm.types.utils import CountContextManagement, TokenCountResponse
API_BASE: Final = "https://my-resource.services.ai.azure.com"
COUNT_URL: Final = f"{API_BASE}/anthropic/v1/messages/count_tokens"
@ -268,3 +271,92 @@ async def test_count_sends_the_tool_change_beta_when_the_history_adds_a_tool(mon
assert result is not None
assert result.total_tokens == 541
assert route.calls.last.request.headers["anthropic-beta"] == "mid-conversation-tool-changes-2026-07-01,token-counting-2024-11-01"
async def _betas_azure_inference_sends(context_management: Mapping[str, object]) -> frozenset[str]:
with respx.mock:
route: Final = respx.post(f"{API_BASE}/anthropic/v1/messages").mock(
return_value=httpx.Response(
200,
json={
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-5-5",
"content": [{"type": "text", "text": "ok"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 12, "output_tokens": 2},
},
)
)
await litellm.anthropic.messages.acreate(
max_tokens=64,
messages=[{"role": "user", "content": "Summarize the thread."}],
model="azure_ai/claude-sonnet-5-5",
api_key="azure-key",
api_base=API_BASE,
context_management=context_management,
)
return frozenset(filter(None, route.calls.last.request.headers.get("anthropic-beta", "").split(",")))
@pytest.mark.asyncio
@pytest.mark.usefixtures("local_beta_headers_config")
@pytest.mark.parametrize("edit_type", ["clear_tool_uses_20250919", "compact_20260112"])
async def test_count_forwards_the_context_edits_with_the_betas_azure_inference_sends(
monkeypatch: pytest.MonkeyPatch, edit_type: str
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
spec: Final = {"edits": [{"type": edit_type}]}
inference_betas: Final = await _betas_azure_inference_sends(spec)
with respx.mock:
route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 88}))
result: Final = await AzureAIAnthropicTokenCounter().count_tokens(
model_to_use="claude-sonnet-5-5",
messages=[{"role": "user", "content": "Summarize the thread."}],
contents=None,
deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}},
request_model="azure_ai/claude-sonnet-5-5",
context_management=CountContextManagement(spec=spec, drop_invalid_edits=False),
)
assert result is not None
assert result.total_tokens == 88
assert json.loads(route.calls.last.request.content)["context_management"] == spec
assert frozenset(route.calls.last.request.headers["anthropic-beta"].split(",")) == inference_betas | {
"token-counting-2024-11-01"
}
class _FetchThreads:
seen: tuple[int, ...] = ()
def __call__(self, request: httpx.Request) -> httpx.Response:
self.seen = (*self.seen, threading.get_ident())
return httpx.Response(200, json={})
@pytest.mark.asyncio
async def test_count_with_an_edit_fetches_the_remote_beta_config_off_the_event_loop(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False)
monkeypatch.setattr(anthropic_beta_headers_manager, "_BETA_HEADERS_CONFIG", None)
fetch_threads: Final = _FetchThreads()
with respx.mock:
respx.get(litellm.anthropic_beta_headers_url).mock(side_effect=fetch_threads)
respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 88}))
result: Final = await AzureAIAnthropicTokenCounter().count_tokens(
model_to_use="claude-sonnet-5-5",
messages=[{"role": "user", "content": "Summarize the thread."}],
contents=None,
deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}},
request_model="azure_ai/claude-sonnet-5-5",
context_management=CountContextManagement(
spec={"edits": [{"type": "clear_tool_uses_20250919"}]}, drop_invalid_edits=False
),
)
assert result is not None
assert len(fetch_threads.seen) == 1
assert threading.get_ident() not in fetch_threads.seen

18
tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py
View file

@ -1,6 +1,8 @@
import base64
import json
import pytest
from litellm.llms.bedrock.count_tokens.transformation import (
DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS,
BedrockCountTokensConfig,
@ -53,6 +55,22 @@ def test_detect_input_type_converse_blocks_route_to_converse():
assert config._detect_input_type(request) == "converse"
@pytest.mark.parametrize("content", [[{"text": "hi"}], "hi"])
@pytest.mark.parametrize(
"context_management", [None, {"edits": [{"type": "clear_tool_uses_20250919"}]}]
)
def test_a_pass_through_body_with_context_management_keeps_its_converse_input(content, context_management):
request = {
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
"messages": [{"role": "user", "content": content}],
"context_management": context_management,
}
transformed = BedrockCountTokensConfig().transform_anthropic_to_bedrock_count_tokens(request)
assert transformed["input"]["converse"]["messages"] == [{"role": "user", "content": [{"text": "hi"}]}]
def test_transform_to_invoke_model_format_base64_encodes_body():
"""The CountTokens API expects invokeModel.body as a base64-encoded blob;
Anthropic Messages bodies additionally need anthropic_version/max_tokens

99
tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py
View file

@ -9,6 +9,7 @@ from pydantic import JsonValue, TypeAdapter
import litellm
from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.types.utils import CountContextManagement
RUNTIME_HOST: Final = "bedrock-runtime.us-east-1.amazonaws.com"
MANTLE_HOST: Final = "bedrock-mantle.us-east-1.api.aws"
@ -144,6 +145,104 @@ async def test_mantle_failure_is_reported_with_its_status() -> None:
assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST)
@pytest.mark.asyncio
async def test_runtime_rejection_of_an_edit_stands_when_mantle_does_not_serve_the_model() -> None:
rejection: Final = "'claude-haiku-4-5-20251001' does not support the 'compact_20260112' context management strategy"
bedrock: Final = _Bedrock(
runtime=(400, {"message": rejection}),
mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}),
)
result: Final = await _counter(bedrock).count_tokens(
model_to_use="us.anthropic.claude-haiku-4-5-20251001-v1:0",
messages=MESSAGES,
contents=None,
deployment=DEPLOYMENT,
request_model="claude-haiku-4-5",
context_management=CountContextManagement(
spec={"edits": [{"type": "compact_20260112"}]}, drop_invalid_edits=False
),
)
assert result is not None
assert result.error is True
assert result.status_code == 400
assert result.tokenizer_type == "bedrock_api"
assert rejection in (result.error_message or "")
assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST)
@pytest.mark.asyncio
async def test_runtime_400_without_an_edit_still_reports_the_mantle_failure() -> None:
bedrock: Final = _Bedrock(
runtime=(400, {"message": "The provided model identifier is invalid."}),
mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}),
)
result: Final = await _counter(bedrock).count_tokens(
model_to_use="us.anthropic.claude-haiku-4-5-20251001-v1:0",
messages=MESSAGES,
contents=None,
deployment=DEPLOYMENT,
request_model="claude-haiku-4-5",
)
assert result is not None
assert result.status_code == 404
assert result.tokenizer_type == "bedrock_mantle_api"
assert "does not exist" in (result.error_message or "")
assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST)
@pytest.mark.asyncio
async def test_plain_text_history_with_an_edit_is_counted_through_invoke_carrying_the_edit() -> None:
bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 57}), mantle=(200, {"input_tokens": 1}))
spec: Final[dict[str, JsonValue]] = {"edits": [{"type": "clear_tool_uses_20250919"}]}
result: Final = await _counter(bedrock).count_tokens(
model_to_use="global.anthropic.claude-opus-4-8",
messages=MESSAGES,
contents=None,
deployment=DEPLOYMENT,
request_model="claude-opus-4-8",
context_management=CountContextManagement(spec=spec, drop_invalid_edits=False),
)
assert result is not None
assert result.total_tokens == 57
assert bedrock.posted_hosts() == (RUNTIME_HOST,)
invoke_body: Final = _invoke_body(bedrock.requests[0])
assert invoke_body["context_management"] == spec
assert invoke_body["anthropic_beta"] == ["context-management-2025-06-27"]
@pytest.mark.asyncio
async def test_an_empty_string_turn_with_an_edit_is_sent_as_an_empty_string_not_an_empty_text_block() -> None:
bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 16}), mantle=(200, {"input_tokens": 1}))
spec: Final[dict[str, JsonValue]] = {"edits": [{"type": "clear_tool_uses_20250919"}]}
result: Final = await _counter(bedrock).count_tokens(
model_to_use="global.anthropic.claude-opus-4-8",
messages=[
{"role": "user", "content": "first"},
{"role": "assistant", "content": ""},
{"role": "user", "content": "third question here"},
],
contents=None,
deployment=DEPLOYMENT,
request_model="claude-opus-4-8",
context_management=CountContextManagement(spec=spec, drop_invalid_edits=False),
)
assert result is not None
assert result.total_tokens == 16
assert _invoke_body(bedrock.requests[0])["messages"] == [
{"role": "user", "content": [{"type": "text", "text": "first"}]},
{"role": "assistant", "content": ""},
{"role": "user", "content": [{"type": "text", "text": "third question here"}]},
]
@pytest.mark.asyncio
async def test_claude_runtime_404_retries_with_system_and_tools_on_mantle() -> None:
bedrock: Final = _Bedrock(runtime=(404, UNSUPPORTED), mantle=(200, {"input_tokens": 2177}))

230
tests/unit/llms/vertex_ai/count_tokens/test_context.py
View file

@ -1,6 +1,5 @@
import asyncio
import base64
import logging
import uuid
from collections.abc import Generator
from contextlib import contextmanager
@ -11,20 +10,22 @@ from typing import Final, Literal
from unittest.mock import patch
import pytest
from pydantic import JsonValue, TypeAdapter
from pydantic import JsonValue
import litellm
from litellm.llms.anthropic.messages.adapters.transformation import (
LiteLLMAnthropicMessagesAdapter,
truncate_tool_name,
)
from litellm.llms.vertex_ai.count_tokens.context import (
from litellm.llms.anthropic.messages.context_management.inference_view import (
CountContextEdits,
InferenceRequestView,
inference_request_view,
)
from litellm.llms.vertex_ai.count_tokens.context import (
VertexGeminiCountContext,
VertexGeminiCountContextError,
gemini_count_contents,
inference_request_view,
unconvertible_messages_error,
vertex_gemini_count_context,
)
@ -424,73 +425,6 @@ async def test_image_that_cannot_be_fetched_returns_an_error_value(
assert isinstance(contents, VertexGeminiCountContextError)
@pytest.mark.parametrize(
"system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system")
)
async def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None:
view: Final = await inference_request_view(
(
{"role": "user", "content": "first question " * 50},
{"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]},
{"role": "user", "content": "latest question"},
),
system,
"anthropic",
_NO_EDITS,
)
assert isinstance(view, InferenceRequestView)
assert view.messages == [{"role": "user", "content": "latest question"}]
assert "Summary of the earlier turns." in str(view.system)
assert "Be brief." in str(view.system)
async def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None:
snippet: Final = "LiteLLM ships a weekly release candidate and a stable line."
view: Final = await inference_request_view(
(
{"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,
}
],
},
],
},
),
None,
"anthropic",
_NO_EDITS,
)
assert isinstance(view, InferenceRequestView)
assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"])
assert [block["type"] for block in assistant_blocks] == ["text"]
assert snippet in assistant_blocks[0]["text"]
async def test_openai_view_is_the_request_unchanged() -> None:
messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},)
view: Final = await inference_request_view(messages, "Be brief.", "openai", _NO_EDITS)
assert isinstance(view, InferenceRequestView)
assert view.messages == messages
assert view.system == "Be brief."
async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> None:
with patch(
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type",
@ -556,138 +490,30 @@ async def test_anthropic_replayed_web_search_result_is_counted_as_its_text(
assert snippet in str(contents)
_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180
_READ_FILE_TOOL: Final[dict[str, object]] = {
"name": "read_file",
"description": "Read a file",
"input_schema": {"type": "object", "properties": {"path": {"type": "string"}}},
}
_CLEAR_EDIT: Final[JsonValue] = {
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": 500},
"keep": {"type": "tool_uses", "value": 1},
}
_CLEAR_OLD_TOOL_RESULTS: Final[JsonValue] = {"edits": [_CLEAR_EDIT]}
def _agent_turns(rounds: int, first_round: int = 0) -> tuple[dict[str, JsonValue], ...]:
return (
*(
turn
for index in range(first_round, first_round + rounds)
for turn in (
{"role": "user", "content": f"Read file {index}."},
{
"role": "assistant",
"content": [
{"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}"}}
],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": f"toolu_{index}", "content": _TOOL_OUTPUT}],
},
)
), # comprehension-ok: flattens each round into its three turns
{"role": "user", "content": "Summarize what you read."},
)
def _tool_outputs_kept(view: InferenceRequestView) -> tuple[bool, ...]:
return tuple(
_TOOL_OUTPUT in str(message["content"]) for message in view.messages if "tool_result" in str(message["content"])
)
def _edits(spec: JsonValue) -> CountContextEdits:
return CountContextEdits(spec=spec, model="gemini/gemini-2.5-flash", tools=(_READ_FILE_TOOL,))
async def test_anthropic_view_clears_the_tool_results_inference_clears() -> None:
view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS))
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (False, False, True)
async def test_anthropic_view_applies_clear_edits_without_calling_the_compaction_summary_model() -> None:
spec: Final[JsonValue] = {
"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50000}}, _CLEAR_EDIT]
}
turns: Final = (
{"role": "user", "content": "first question " * 30000},
{"role": "assistant", "content": "first answer"},
*_agent_turns(3),
)
with (
patch(
"litellm.llms.anthropic.messages.context_management.editors.compact._read_summary_model_setting",
return_value="gemini/gemini-2.5-flash",
),
patch.object(litellm, "acompletion", side_effect=AssertionError("count_tokens must not summarize")) as summary,
):
view: Final = await inference_request_view(turns, None, "anthropic", _edits(spec))
assert summary.call_count == 0
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (False, False, True)
async def test_anthropic_view_counts_the_unedited_body_without_an_error_log_when_an_edit_fails(
caplog: pytest.LogCaptureFixture,
) -> None:
with (
patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")),
caplog.at_level(logging.DEBUG, logger="LiteLLM"),
):
view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS))
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (True, True, True)
assert any("skipping edits" in record.getMessage() for record in caplog.records)
assert not any(record.levelno >= logging.ERROR for record in caplog.records)
async def test_anthropic_view_clears_only_the_turns_after_a_compaction_block_when_clear_runs_before_compact() -> None:
clear_before_compact: Final[JsonValue] = {
"edits": [
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "tool_uses", "value": 3},
"keep": {"type": "tool_uses", "value": 1},
},
{"type": "compact_20260112"},
]
}
turns: Final = (
*_agent_turns(3)[:-1],
{"role": "assistant", "content": [{"type": "compaction", "content": "Read files 0 to 2."}]},
*_agent_turns(2, first_round=3),
)
with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")):
view: Final = await inference_request_view(turns, None, "anthropic", _edits(clear_before_compact))
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (True, True)
assert "Read files 0 to 2." in str(view.system)
async def test_anthropic_view_rejects_a_compact_trigger_inference_rejects() -> None:
low_trigger: Final[JsonValue] = {
"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}]
}
view: Final = await inference_request_view(_agent_turns(1), None, "anthropic", _edits(low_trigger))
assert isinstance(view, VertexGeminiCountContextError)
assert view.status_code == 400
async def test_anthropic_view_names_the_block_gemini_cannot_count_when_an_edit_trigger_cannot_tokenize_it() -> None:
clear_old_tool_results: Final = CountContextEdits(
spec={
"edits": [
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": 500},
"keep": {"type": "tool_uses", "value": 1},
}
]
},
model="gemini/gemini-2.5-flash",
tools=({"name": "read_file", "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}},),
)
turns: Final = (
*_agent_turns(1)[:-1],
{"role": "user", "content": "Read file 0."},
{
"role": "assistant",
"content": [{"type": "tool_use", "id": "toolu_0", "name": "read_file", "input": {"path": "0"}}],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "toolu_0", "content": "line of a long file. " * 180}],
},
{
"role": "user",
"content": [
@ -701,7 +527,7 @@ async def test_anthropic_view_names_the_block_gemini_cannot_count_when_an_edit_t
},
)
view: Final = await inference_request_view(turns, None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS))
view: Final = await inference_request_view(turns, None, "anthropic", clear_old_tool_results)
assert isinstance(view, InferenceRequestView)
blocks_error: Final = unconvertible_messages_error(view.messages, "anthropic")

99
tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py
View file

@ -1,16 +1,19 @@
import copy
import threading
import time
from collections.abc import Iterator
from contextlib import contextmanager
from enum import Enum
from typing import Final
from unittest.mock import AsyncMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import respx
from pydantic import JsonValue, TypeAdapter
import litellm
from litellm import anthropic_beta_headers_manager
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler
@ -1826,6 +1829,100 @@ async def test_vertex_ai_token_counter_forwards_system_and_tools_to_partner_requ
assert without_optional_fields.total_tokens == 37
@pytest.mark.asyncio
@pytest.mark.parametrize(
("edit_type", "expected_beta"),
[("clear_tool_uses_20250919", "context-management-2025-06-27"), ("compact_20260112", "compact-2026-01-12")],
)
async def test_vertex_claude_count_forwards_the_context_edits_and_beta_vertex_inference_sends(
edit_type: str, expected_beta: str, monkeypatch: pytest.MonkeyPatch
) -> None:
from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler
class FakeResponse:
status_code = 200
def json(self) -> dict[str, int]:
return {"input_tokens": 21}
class FakeHttpClient:
posted: tuple[tuple[dict[str, str], dict[str, object]], ...] = ()
async def post(
self, url: str, headers: dict[str, str], json: dict[str, object], timeout: float
) -> FakeResponse:
self.posted = (*self.posted, (headers, json))
return FakeResponse()
fake_http_client: Final = FakeHttpClient()
spec: Final[dict[str, JsonValue]] = {"edits": [{"type": edit_type}]}
application_default_credentials: Final = MagicMock(token="fake-token", valid=True, expiry=None)
monkeypatch.setattr("google.auth.default", lambda *args, **kwargs: (application_default_credentials, "test-project"))
with patch.object(handler, "get_async_httpx_client", return_value=fake_http_client):
result: Final = await VertexAITokenCounter().count_tokens(
model_to_use="claude-opus-4-8",
messages=[{"role": "user", "content": "Hello"}],
contents=None,
deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-east5"}},
messages_format="anthropic",
context_management=CountContextManagement(spec=spec, drop_invalid_edits=False),
)
assert result is not None
assert result.total_tokens == 21
((headers, body),) = fake_http_client.posted
assert body["context_management"] == spec
assert headers["anthropic-beta"] == expected_beta
@pytest.mark.asyncio
async def test_vertex_claude_count_with_an_edit_fetches_the_remote_beta_config_off_the_event_loop(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler
class FakeResponse:
status_code = 200
def json(self) -> dict[str, int]:
return {"input_tokens": 21}
class FakeHttpClient:
async def post(
self, url: str, headers: dict[str, str], json: dict[str, object], timeout: float
) -> FakeResponse:
return FakeResponse()
class FetchThreads:
seen: tuple[int, ...] = ()
def __call__(self, request: httpx.Request) -> httpx.Response:
self.seen = (*self.seen, threading.get_ident())
return httpx.Response(200, json={})
monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False)
monkeypatch.setattr(anthropic_beta_headers_manager, "_BETA_HEADERS_CONFIG", None)
fetch_threads: Final = FetchThreads()
application_default_credentials: Final = MagicMock(token="fake-token", valid=True, expiry=None)
monkeypatch.setattr("google.auth.default", lambda *args, **kwargs: (application_default_credentials, "test-project"))
with patch.object(handler, "get_async_httpx_client", return_value=FakeHttpClient()), respx.mock:
respx.get(litellm.anthropic_beta_headers_url).mock(side_effect=fetch_threads)
result: Final = await VertexAITokenCounter().count_tokens(
model_to_use="claude-opus-4-8",
messages=[{"role": "user", "content": "Hello"}],
contents=None,
deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-east5"}},
messages_format="anthropic",
context_management=CountContextManagement(
spec={"edits": [{"type": "clear_tool_uses_20250919"}]}, drop_invalid_edits=False
),
)
assert result is not None
assert len(fetch_threads.seen) == 1
assert threading.get_ident() not in fetch_threads.seen
_PARTNER_CHAT_TOOL_TURNS: Final[list[dict[str, object]]] = [
{"role": "user", "content": "Weather in SF?"},
{

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

@ -25,7 +25,7 @@ import pytest
import respx
import tiktoken
from dotenv import load_dotenv
from pydantic import BaseModel, ConfigDict, JsonValue
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter
from tokenizers import Tokenizer
from tokenizers.models import WordLevel
from tokenizers.pre_tokenizers import Whitespace
@ -82,6 +82,7 @@ from litellm.proxy.proxy_server import (
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
from litellm.types.utils import (
ADDRESSED_RESPONSE_ID_FIELD,
@ -1741,7 +1742,10 @@ _OPUS_IN_US_EAST_1: Final = TokenCountModel(model="us.anthropic.claude-opus-5-5"
async def _ladder(
counter: _ScriptedCounter, settings: TokenCountSettings, local_estimate_allowed: bool = True
counter: _ScriptedCounter,
settings: TokenCountSettings,
local_estimate_allowed: bool = True,
forwarded_edits: bool = False,
) -> TokenCountResponse | None:
return await _provider_token_count(
count_with=partial(
@ -1759,6 +1763,7 @@ async def _ladder(
request_model=_COUNT_REQUEST.model,
settings=settings,
local_estimate_allowed=local_estimate_allowed,
forwarded_edits=forwarded_edits,
)
@ -1986,6 +1991,26 @@ async def test_configured_token_count_model_error_stands_even_with_framing(
assert raised.value.message.count("counts tokens with us.anthropic.claude-opus-5-5 in us-east-1") == region_hints
@pytest.mark.parametrize(("status_code", "region_hints"), [(400, 0), (403, 1)])
@pytest.mark.asyncio
async def test_a_pinned_token_count_model_rejecting_forwarded_edits_answers_without_the_region_hint(
status_code: int, region_hints: int
):
counter: Final = _ScriptedCounter(
{"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)}
)
with pytest.raises(ProxyException) as raised:
await _ladder(
counter,
TokenCountSettings(framing=None, count_model=replace(_OPUS_IN_US_EAST_1, configured=True)),
forwarded_edits=True,
)
assert raised.value.code == str(status_code)
assert raised.value.message.count("counts tokens with us.anthropic.claude-opus-5-5 in us-east-1") == region_hints
@pytest.mark.parametrize(
("count_model", "region_hints"),
[(_OPUS_IN_US_EAST_1, frozenset({True})), (replace(_OPUS_IN_US_EAST_1, region=None), frozenset())],
@ -3172,15 +3197,14 @@ _READ_FILE_TOOL: Final[dict[str, JsonValue]] = {
"name": "read_file",
"input_schema": {"type": "object", "properties": {"path": {"type": "string"}}},
}
_CLEAR_TOOL_USES: Final[dict[str, JsonValue]] = {
"edits": [
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": 500},
"keep": {"type": "tool_uses", "value": 1},
}
]
}
_CLEAR_TOOL_USES_EDITS: Final[list[JsonValue]] = [
{
"type": "clear_tool_uses_20250919",
"trigger": {"type": "input_tokens", "value": 500},
"keep": {"type": "tool_uses", "value": 1},
}
]
_CLEAR_TOOL_USES: Final[dict[str, JsonValue]] = {"edits": _CLEAR_TOOL_USES_EDITS}
_INVALID_COMPACT_EDIT: Final[dict[str, JsonValue]] = {
"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}]
}
@ -3516,6 +3540,28 @@ async def test_gemini_messages_count_tokens_counts_the_unedited_body_when_drop_p
assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == 3
@pytest.mark.asyncio
@pytest.mark.parametrize(
("router_defaults", "request_options", "expected_status"),
(({}, {"drop_params": True}, None), ({"drop_params": True}, {}, None), ({}, {}, 400)),
ids=("request-drop-params", "router-default-drop-params", "no-drop-params"),
)
async def test_messages_count_tokens_for_a_model_with_no_deployment_drops_an_invalid_edit_only_under_drop_params(
monkeypatch: pytest.MonkeyPatch,
router_defaults: dict[str, JsonValue],
request_options: dict[str, JsonValue],
expected_status: int | None,
) -> None:
monkeypatch.setenv("GEMINI_API_KEY", "test-key")
with patch.object(litellm, "drop_params", False):
response, _ = await _count_through_the_messages_endpoint(
Router(model_list=[_gemini_deployment("gemini-agent", {})], default_litellm_params=router_defaults),
{**request_options, "model": "gemini/gemini-3.8-pro", "context_management": _INVALID_COMPACT_EDIT},
)
assert (response.status_code if isinstance(response, JSONResponse) else None) == expected_status
@pytest.mark.asyncio
async def test_gemini_messages_count_tokens_rejects_an_invalid_edit_without_drop_params() -> None:
with patch.object(litellm, "drop_params", False):
@ -3529,6 +3575,576 @@ async def test_gemini_messages_count_tokens_rejects_an_invalid_edit_without_drop
assert sent_bodies == ()
_MISTRAL_MODEL: Final = "mistral/mistral-large-latest"
def _mistral_router(
litellm_params: Mapping[str, JsonValue] = MappingProxyType({}),
router_defaults: Mapping[str, JsonValue] = MappingProxyType({}),
) -> Router:
return Router(
model_list=[
{
"model_name": "mistral-agent",
"litellm_params": {"model": _MISTRAL_MODEL, "api_key": "test-key", **litellm_params},
}
],
default_litellm_params=dict(router_defaults),
)
async def _messages_endpoint_count(router: Router, request_body: Mapping[str, JsonValue]) -> object:
import litellm.proxy.anthropic_endpoints.endpoints as anthropic_endpoints
import litellm.proxy.proxy_server as proxy_server
mock_request: Final = MagicMock(spec=Request)
mock_request.headers = {}
with (
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())
class _MessagesCountReply(BaseModel):
model_config = ConfigDict(frozen=True)
input_tokens: int
async def _input_tokens(router: Router, request_body: Mapping[str, JsonValue]) -> int:
return _MessagesCountReply.model_validate(await _messages_endpoint_count(router, request_body)).input_tokens
async def _messages_the_inference_polyfill_sends(model: str) -> list[dict[str, JsonValue]]:
from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request
prepared: Final = await prepare_context_managed_request(
model=model,
messages=[dict(message) for message in _agent_turns_with_tool_results(3)],
tools=[dict(_READ_FILE_TOOL)],
system=None,
context_management_spec=dict(_CLEAR_TOOL_USES),
litellm_metadata=None,
additional_drop_params=None,
llm_router=None,
)
assert prepared is not None
return TypeAdapter(list[dict[str, JsonValue]]).validate_python(prepared.messages)
@pytest.mark.asyncio
async def test_messages_count_tokens_counts_the_body_the_inference_polyfill_sends_on_a_bridged_provider() -> None:
router: Final = _mistral_router()
unedited_body: Final[dict[str, JsonValue]] = {
"model": "mistral-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
}
edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES})
unedited: Final = await _input_tokens(router, unedited_body)
polyfilled: Final = await _input_tokens(
router, {**unedited_body, "messages": await _messages_the_inference_polyfill_sends(_MISTRAL_MODEL)}
)
assert edited < unedited
assert edited == polyfilled
@pytest.mark.asyncio
@pytest.mark.parametrize(
("deployment_params", "router_defaults", "request_options"),
(
({"additional_drop_params": ["context_management"]}, {}, {}),
({}, {"additional_drop_params": ["context_management"]}, {}),
({}, {}, {"additional_drop_params": ["context_management"]}),
),
ids=("deployment-drops", "router-default-drops", "request-drops"),
)
async def test_messages_count_tokens_on_a_bridged_provider_keeps_the_unedited_count_when_context_management_is_dropped(
deployment_params: dict[str, JsonValue],
router_defaults: dict[str, JsonValue],
request_options: dict[str, JsonValue],
) -> None:
router: Final = _mistral_router(deployment_params, router_defaults)
unedited_body: Final[dict[str, JsonValue]] = {
"model": "mistral-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
**request_options,
}
edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES})
assert edited == await _input_tokens(router, unedited_body)
@pytest.mark.asyncio
async def test_messages_count_tokens_on_a_bridged_provider_rejects_an_invalid_edit_without_drop_params() -> None:
with patch.object(litellm, "drop_params", False):
response: Final = await _messages_endpoint_count(
_mistral_router(),
{
"model": "mistral-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"context_management": _INVALID_COMPACT_EDIT,
},
)
assert isinstance(response, JSONResponse)
assert response.status_code == 400
@pytest.mark.asyncio
async def test_messages_count_tokens_on_a_bridged_provider_counts_without_an_invalid_edit_under_drop_params() -> None:
router: Final = _mistral_router()
unedited_body: Final[dict[str, JsonValue]] = {
"model": "mistral-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"drop_params": True,
}
with patch.object(litellm, "drop_params", False):
edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _INVALID_COMPACT_EDIT})
unedited: Final = await _input_tokens(router, unedited_body)
assert edited == unedited
def _turns_compacted_by_the_client() -> list[dict[str, JsonValue]]:
return [
*_agent_turns_with_tool_results(3)[:-1],
{"role": "assistant", "content": [{"type": "compaction", "content": "The agent read files 0, 1 and 2."}]},
{"role": "user", "content": "Summarize what you read."},
]
@pytest.mark.asyncio
async def test_messages_count_tokens_on_a_bridged_provider_counts_only_what_follows_a_client_compaction_block() -> None:
from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request
router: Final = _mistral_router()
sent: Final = await prepare_context_managed_request(
model=_MISTRAL_MODEL,
messages=[dict(message) for message in _turns_compacted_by_the_client()],
tools=[dict(_READ_FILE_TOOL)],
system=None,
context_management_spec=None,
litellm_metadata=None,
additional_drop_params=None,
llm_router=None,
)
assert sent is not None
compacted: Final = await _input_tokens(
router, {"model": "mistral-agent", "messages": _turns_compacted_by_the_client(), "tools": [_READ_FILE_TOOL]}
)
uncompacted: Final = await _input_tokens(
router, {"model": "mistral-agent", "messages": _agent_turns_with_tool_results(3), "tools": [_READ_FILE_TOOL]}
)
sent_body_count: Final = await _input_tokens(
router,
{
"model": "mistral-agent",
"messages": TypeAdapter(list[dict[str, JsonValue]]).validate_python(sent.messages),
"system": TypeAdapter(JsonValue).validate_python(sent.system),
"tools": [_READ_FILE_TOOL],
},
)
assert compacted < uncompacted
assert compacted == sent_body_count
@pytest.mark.asyncio
async def test_messages_count_tokens_applies_the_edits_a_responses_api_bridge_polyfills_for_a_non_openai_provider() -> (
None
):
perplexity_model: Final = "perplexity/perplexity/sonar"
router: Final = Router(
model_list=[{"model_name": "sonar-agent", "litellm_params": {"model": perplexity_model, "api_key": "test-key"}}]
)
unedited_body: Final[dict[str, JsonValue]] = {
"model": "sonar-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
}
edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES})
unedited: Final = await _input_tokens(router, unedited_body)
polyfilled: Final = await _input_tokens(
router, {**unedited_body, "messages": await _messages_the_inference_polyfill_sends(perplexity_model)}
)
assert edited < unedited
assert edited == polyfilled
@pytest.mark.asyncio
async def test_messages_count_tokens_on_a_responses_api_bridge_counts_the_unedited_body_when_an_edit_fails() -> None:
from litellm.llms.anthropic.messages.responses_adapters import handler as responses_adapters_handler
router: Final = Router(
model_list=[
{
"model_name": "sonar-agent",
"litellm_params": {"model": "perplexity/perplexity/sonar", "api_key": "test-key"},
}
]
)
unedited_body: Final[dict[str, JsonValue]] = {
"model": "sonar-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
}
with patch.object(
responses_adapters_handler, "apply_context_management", AsyncMock(side_effect=RuntimeError("editor bug"))
):
edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES})
assert edited == await _input_tokens(router, unedited_body)
@pytest.mark.asyncio
async def test_messages_count_tokens_on_a_responses_api_bridge_rejects_an_invalid_edit_without_drop_params() -> None:
router: Final = Router(
model_list=[
{
"model_name": "sonar-agent",
"litellm_params": {"model": "perplexity/perplexity/sonar", "api_key": "test-key"},
}
]
)
invalid_clear_edit: Final[dict[str, JsonValue]] = {
"edits": [{"type": "clear_tool_uses_20250919", "clear_at_least": "all of them"}]
}
with patch.object(litellm, "drop_params", False):
response: Final = await _messages_endpoint_count(
router,
{
"model": "sonar-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"context_management": invalid_clear_edit,
},
)
assert isinstance(response, JSONResponse)
assert response.status_code == 400
message: Final = json.loads(bytes(response.body))["error"]["message"]
assert "context_management.edits.0.clear_tool_uses_20250919.clear_at_least" in message
class _CountTokensPeer:
def __init__(self, reply: Mapping[str, JsonValue]) -> None:
self._reply: Final = reply
self.requests: tuple[httpx.Request, ...] = ()
def __call__(self, request: httpx.Request) -> httpx.Response:
self.requests = (*self.requests, request)
return httpx.Response(200, json=dict(self._reply), request=request)
class _BedrockInvokeModel(BaseModel):
model_config = ConfigDict(frozen=True)
body: str
class _BedrockCountTokensInput(BaseModel):
model_config = ConfigDict(frozen=True)
invoke_model: _BedrockInvokeModel = Field(alias="invokeModel")
class _BedrockCountTokensBody(BaseModel):
model_config = ConfigDict(frozen=True)
input: _BedrockCountTokensInput
def _bedrock_counted_request(request_body: bytes) -> str:
invoke_model: Final = _BedrockCountTokensBody.model_validate_json(request_body).input.invoke_model
return base64.b64decode(invoke_model.body).decode()
async def _count_on_a_bridged_provider_count_api(
request_body: Mapping[str, JsonValue],
monkeypatch: pytest.MonkeyPatch,
bedrock_model: str = "bedrock/us.amazon.nova-2-lite-v1:0",
peer_reply: Mapping[str, JsonValue] = MappingProxyType({"inputTokens": 57}),
model_info: Mapping[str, JsonValue] = MappingProxyType({}),
) -> tuple[object, tuple[httpx.Request, ...]]:
from litellm.llms.bedrock.count_tokens import handler as bedrock_count_tokens_handler
from litellm.llms.bedrock.count_tokens import mantle_handler as bedrock_mantle_count_tokens_handler
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
peer: Final = _CountTokensPeer(peer_reply)
router: Final = Router(
model_list=[
{
"model_name": "nova-agent",
"litellm_params": {
"model": bedrock_model,
"aws_access_key_id": "AKIATESTACCESSKEY",
"aws_secret_access_key": "test-secret",
"aws_region_name": "us-east-1",
},
"model_info": dict(model_info),
}
]
)
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer))
with (
patch.object(bedrock_count_tokens_handler, "get_async_httpx_client", return_value=client),
patch.object(bedrock_mantle_count_tokens_handler, "get_async_httpx_client", return_value=client),
):
response: Final = await _messages_endpoint_count(router, {**request_body, "model": "nova-agent"})
return response, peer.requests
@pytest.mark.asyncio
async def test_messages_count_tokens_sends_the_edited_messages_to_a_bridged_provider_count_api(
monkeypatch: pytest.MonkeyPatch,
) -> None:
response, sent_bodies = await _count_on_a_bridged_provider_count_api(
{
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"context_management": _CLEAR_TOOL_USES,
},
monkeypatch,
)
assert response == {"input_tokens": 57, "litellm_count_source": "provider"}
assert len(sent_bodies) == 1
assert _bedrock_counted_request(sent_bodies[0].content).count(_GEMINI_TOOL_OUTPUT) == 1
@pytest.mark.asyncio
async def test_messages_count_tokens_sends_a_body_with_only_empty_text_to_the_bridged_provider_count_api(
monkeypatch: pytest.MonkeyPatch,
) -> None:
empty_turn: Final[dict[str, JsonValue]] = {"role": "user", "content": [{"type": "text", "text": ""}]}
response, sent_bodies = await _count_on_a_bridged_provider_count_api({"messages": [empty_turn]}, monkeypatch)
assert response == {"input_tokens": 57, "litellm_count_source": "provider"}
assert len(sent_bodies) == 1
_BEDROCK_CLAUDE_MODEL: Final = "bedrock/us.anthropic.claude-sonnet-5-5"
_CLEAR_THINKING_EDIT: Final[dict[str, JsonValue]] = {"type": "clear_thinking_20251015"}
_COUNTED_ON_MANTLE: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"token_count_model": {"model": "us.anthropic.claude-sonnet-5-5", "litellm_provider": "bedrock_mantle"}}
)
_COUNTED_ON_RUNTIME: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"token_count_model": {"model": "us.anthropic.claude-sonnet-5-5", "litellm_provider": "bedrock"}}
)
def _bedrock_claude_count_request(edits: list[JsonValue]) -> Mapping[str, JsonValue]:
return MappingProxyType(
{
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"context_management": {"edits": edits},
}
)
@pytest.mark.asyncio
async def test_messages_count_tokens_forwards_bedrock_supported_edits_and_their_beta_to_the_mantle_count_api(
monkeypatch: pytest.MonkeyPatch,
) -> None:
response, sent = await _count_on_a_bridged_provider_count_api(
_bedrock_claude_count_request([*_CLEAR_TOOL_USES_EDITS, _CLEAR_THINKING_EDIT]),
monkeypatch,
_BEDROCK_CLAUDE_MODEL,
MappingProxyType({"input_tokens": 57, "context_management": {"original_input_tokens": 4100}}),
_COUNTED_ON_MANTLE,
)
assert isinstance(response, dict)
assert (response["input_tokens"], response["context_management"]) == (57, {"original_input_tokens": 4100})
assert len(sent) == 1
assert json.loads(sent[0].content)["context_management"] == _CLEAR_TOOL_USES
assert sent[0].headers["anthropic-beta"] == ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value
@pytest.mark.asyncio
async def test_messages_count_tokens_forwards_bedrock_supported_edits_and_their_beta_to_the_runtime_count_api(
monkeypatch: pytest.MonkeyPatch,
) -> None:
response, sent = await _count_on_a_bridged_provider_count_api(
_bedrock_claude_count_request([*_CLEAR_TOOL_USES_EDITS, _CLEAR_THINKING_EDIT]),
monkeypatch,
_BEDROCK_CLAUDE_MODEL,
model_info=_COUNTED_ON_RUNTIME,
)
assert isinstance(response, dict)
assert response["input_tokens"] == 57
assert len(sent) == 1
invoke_body: Final = json.loads(_bedrock_counted_request(sent[0].content))
assert (invoke_body["context_management"], invoke_body["anthropic_beta"]) == (
_CLEAR_TOOL_USES,
[ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value],
)
@pytest.mark.asyncio
async def test_messages_count_tokens_sends_no_context_management_to_bedrock_when_no_edit_is_bedrock_supported(
monkeypatch: pytest.MonkeyPatch,
) -> None:
response, sent = await _count_on_a_bridged_provider_count_api(
_bedrock_claude_count_request([_CLEAR_THINKING_EDIT]),
monkeypatch,
_BEDROCK_CLAUDE_MODEL,
MappingProxyType({"input_tokens": 57}),
_COUNTED_ON_MANTLE,
)
assert isinstance(response, dict)
assert response["input_tokens"] == 57
assert len(sent) == 1
assert "context_management" not in json.loads(sent[0].content)
assert "anthropic-beta" not in sent[0].headers
@pytest.mark.asyncio
async def test_messages_count_tokens_rejects_a_non_string_tool_result_text_before_a_bridged_provider_applies_edits(
monkeypatch: pytest.MonkeyPatch,
) -> None:
response, sent_bodies = await _count_on_a_bridged_provider_count_api(
{"messages": _build_log_messages(5), "context_management": _CLEAR_TOOL_USES}, monkeypatch
)
assert isinstance(response, JSONResponse)
assert response.status_code == 400
assert "tool_result toolu_log content.1.text" in json.loads(bytes(response.body))["error"]["message"]
assert sent_bodies == ()
@pytest.mark.asyncio
async def test_messages_count_tokens_applies_a_responses_api_deployment_edits_once() -> None:
peer: Final = _CountTokensPeer({"object": "response.input_tokens", "input_tokens": 71})
router: Final = Router(
model_list=[{"model_name": "gpt-agent", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "test-key"}}]
)
with patch.object(
openai_count_tokens_http,
"get_async_httpx_client",
return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)),
):
response: Final = await _messages_endpoint_count(
router,
{
"model": "gpt-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"context_management": _CLEAR_TOOL_USES,
},
)
assert response == {"input_tokens": 71, "litellm_count_source": "provider"}
assert len(peer.requests) == 1
assert peer.requests[0].content.decode().count(_GEMINI_TOOL_OUTPUT) == 1
_DATABRICKS_CLAUDE_MODEL: Final = "databricks/databricks-claude-sonnet-4-6"
def _databricks_claude_router() -> Router:
return Router(
model_list=[
{
"model_name": "databricks-claude-agent",
"litellm_params": {
"model": _DATABRICKS_CLAUDE_MODEL,
"api_key": "test-key",
"api_base": "https://example.cloud.databricks.com/serving-endpoints",
},
}
]
)
@pytest.mark.asyncio
async def test_messages_count_tokens_counts_the_edited_body_on_a_native_messages_deployment_without_a_count_api() -> (
None
):
router: Final = _databricks_claude_router()
unedited_body: Final[dict[str, JsonValue]] = {
"model": "databricks-claude-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
}
edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES})
unedited: Final = await _input_tokens(router, unedited_body)
polyfilled: Final = await _input_tokens(
router, {**unedited_body, "messages": await _messages_the_inference_polyfill_sends(_DATABRICKS_CLAUDE_MODEL)}
)
assert edited < unedited
assert edited == polyfilled
@pytest.mark.asyncio
@pytest.mark.parametrize("drop_params", [False, True])
async def test_messages_count_tokens_on_a_native_messages_deployment_without_a_count_api_rejects_an_invalid_edit(
drop_params: bool,
) -> None:
with patch.object(litellm, "drop_params", drop_params):
response: Final = await _messages_endpoint_count(
_databricks_claude_router(),
{
"model": "databricks-claude-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"context_management": _INVALID_COMPACT_EDIT,
},
)
assert isinstance(response, JSONResponse)
assert response.status_code == 400
@pytest.mark.asyncio
async def test_messages_count_tokens_leaves_the_edits_to_a_native_messages_deployment() -> None:
peer: Final = _CountTokensPeer({"input_tokens": 63})
router: Final = Router(
model_list=[
{
"model_name": "claude-agent",
"litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": "test-key"},
}
]
)
with patch.object(
anthropic_count_tokens_handler,
"get_async_httpx_client",
return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)),
):
response: Final = await _messages_endpoint_count(
router,
{
"model": "claude-agent",
"messages": _agent_turns_with_tool_results(3),
"tools": [_READ_FILE_TOOL],
"context_management": _CLEAR_TOOL_USES,
},
)
assert response == {"input_tokens": 63, "litellm_count_source": "provider"}
assert len(peer.requests) == 1
assert peer.requests[0].content.decode().count(_GEMINI_TOOL_OUTPUT) == 3
@pytest.mark.asyncio
@pytest.mark.parametrize(
"text", ["hello there", "The quick brown fox jumps over the lazy dog. " * 50], ids=["short", "long"]