scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-31 06:56:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 245

Merge 62e8e7ac3f into afec145edc

Browse source
This commit is contained in:
tessellark 2026-10-31 06:04:57 +00:00 • committed by GitHub
commit ae6cfe01a0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 1172 additions and 380 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)

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"

56
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,14 @@ def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model:
)
def _mantle_retry_outcome(runtime: CountTokensOutcome, mantle: CountTokensOutcome) -> CountTokensOutcome:
match runtime, mantle:
case CountTokensFailure(status_code=400), CountTokensFailure(status_code=404):
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 +98,26 @@ def _request_data(
}
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 +200,17 @@ 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
)
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,
}
mantle_request: Final = runtime_request if translated is None else {**runtime_request, **translated.options}
primary: Final = (
@ -184,8 +223,11 @@ 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
),
)
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):

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

@ -1565,12 +1565,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 +1587,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,

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,

153
litellm/proxy/proxy_server.py
View file

@ -196,6 +196,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,
@ -386,6 +387,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,
@ -395,6 +402,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
@ -14885,7 +14893,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,
@ -15240,6 +15248,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)
@ -15324,6 +15335,87 @@ 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:
return await inference_request_view(
messages,
system,
messages_format,
count_context_edits(edits, tools, deployment.get("litellm_params", {}), fallback_model),
)
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:
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 {}
@ -15517,6 +15609,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)
@ -15556,16 +15682,16 @@ 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,
@ -15577,9 +15703,8 @@ async def _count_request_tokens(
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:
@ -15597,11 +15722,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(inference_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 inference_messages),
)
counted_messages: Final = (
typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages)
@ -15612,7 +15737,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 edited_request is None
else request.model_copy(update={"messages": inference_messages, "system": inference_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

28
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,33 @@ 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_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")

493
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,
@ -3172,15 +3173,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}}]
}
@ -3529,6 +3529,479 @@ 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_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
_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
@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"]