diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 55eaa0d2ac6..ecf050d185e 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -11,7 +11,7 @@ from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests anthropic_count_tokens_handler: Final = AnthropicCountTokensHandler() @@ -37,6 +37,7 @@ class AnthropicTokenCounter(BaseTokenCounter): system: JsonValue = None, messages_format: TokenCountMessagesFormat = "anthropic", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: """ Count tokens using Anthropic's CountTokens API. diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 60c63964492..7b8c2550838 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -139,7 +139,7 @@ def request_safeguards_evaluator( ) -async def _prepare_context_managed_request( +async def prepare_context_managed_request( *, model: str, messages: _AnthropicMessages, @@ -152,6 +152,8 @@ async def _prepare_context_managed_request( user_api_key_auth: "UserAPIKeyAuth | None" = None, sub_call_logging_kwargs: ProxySubCallLoggingKwargs | None = None, parent_trace_id: str | None = None, + count_only: bool = False, + drop_invalid_edits: bool = False, ) -> PolyfillResult | None: """Apply client compaction history, then optional context_management polyfill.""" from litellm.llms.anthropic.pass_through.context_management.editors.compact import ( @@ -193,6 +195,8 @@ async def _prepare_context_managed_request( user_api_key_auth=user_api_key_auth, sub_call_logging_kwargs=sub_call_logging_kwargs, parent_trace_id=parent_trace_id, + count_only=count_only, + drop_invalid_edits=drop_invalid_edits, ) if polyfill_result is not None: @@ -266,16 +270,14 @@ def _spec_has_non_compact_edits( return any(isinstance(edit.get("type"), str) and edit.get("type") != COMPACT_EDIT_TYPE for edit in edits) -def _context_management_explicitly_dropped(additional_drop_params: list[str] | None) -> bool: +def context_management_explicitly_dropped(additional_drop_params: object) -> bool: """True when the caller opted out of context_management via ``additional_drop_params``. ``drop_params`` deliberately does NOT gate the polyfill: ``context_management`` is a LiteLLM-supported param (native on Anthropic, polyfilled elsewhere), and ``drop_params`` only exists to drop genuinely unsupported params. """ - if not isinstance(additional_drop_params, list): - return False - return "context_management" in additional_drop_params + return isinstance(additional_drop_params, list) and "context_management" in additional_drop_params def _normalize_spec_edits( @@ -291,7 +293,7 @@ def _normalize_spec_edits( if not context_management_spec: return None - if _context_management_explicitly_dropped(additional_drop_params): + if context_management_explicitly_dropped(additional_drop_params): return None from litellm.llms.anthropic.pass_through.context_management.dispatcher import ( @@ -317,20 +319,23 @@ async def _run_polyfill_if_enabled( user_api_key_auth: "UserAPIKeyAuth | None" = None, sub_call_logging_kwargs: ProxySubCallLoggingKwargs | None = None, parent_trace_id: str | None = None, + count_only: bool = False, + drop_invalid_edits: bool = False, ) -> PolyfillResult | None: """Run the async context_management polyfill if a spec is present. Returns ``None`` when the spec is empty or ``context_management`` is listed in ``additional_drop_params`` (the explicit opt-out; ``drop_params`` - does not disable the polyfill because context_management is a supported + does not disable a valid spec because context_management is a supported param). Raises ``AnthropicContextManagementError`` so the /v1/messages - endpoint can emit an Anthropic-format 400. All other exceptions are + endpoint can emit an Anthropic-format 400, unless ``drop_invalid_edits`` + (``drop_params``) drops the invalid spec. All other exceptions are best-effort swallowed (matches v0 behavior). """ if not context_management_spec: return None - if _context_management_explicitly_dropped(additional_drop_params): + if context_management_explicitly_dropped(additional_drop_params): return None try: @@ -345,13 +350,17 @@ async def _run_polyfill_if_enabled( user_api_key_auth=user_api_key_auth, sub_call_logging_kwargs=sub_call_logging_kwargs, parent_trace_id=parent_trace_id, + count_only=count_only, ) except AnthropicContextManagementError: + if drop_invalid_edits: + return None # Surface validation errors so the endpoint can emit an Anthropic-format # 400. Other exception types fall into the best-effort branch below. raise except Exception as e: - verbose_logger.exception("context_management polyfill: skipping edits due to error: %s", e) + log_failure: Final = verbose_logger.debug if count_only else verbose_logger.exception + log_failure("context_management polyfill: skipping edits due to error: %s", e) # Best-effort swallow is only safe for compact-only specs, where the # caller's compaction-block-slicing safety net produces a correct # (if degraded) result. When the spec also requested non-compact @@ -709,6 +718,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: """Handle non-Anthropic models asynchronously using the adapter""" context_management: Final = kwargs.pop("context_management", None) additional_drop_params: Final[list[str] | None] = kwargs.get("additional_drop_params", None) + drop_invalid_edits: Final = litellm.drop_params is True or ( + "drop_params" in kwargs and kwargs["drop_params"] is True + ) litellm_router: Final[Router | None] = _router_or_proxy_fallback(kwargs.pop("litellm_router", None)) proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs) @@ -716,7 +728,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: sub_call_logging_kwargs: Final = proxy_sub_call_logging_kwargs(request_kwargs) sub_call_trace_id: Final = proxy_sub_call_trace_id(request_kwargs) - polyfill_result: Final = await _prepare_context_managed_request( + polyfill_result: Final = await prepare_context_managed_request( model=model, messages=messages, tools=tools, @@ -728,6 +740,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: user_api_key_auth=user_api_key_auth, sub_call_logging_kwargs=sub_call_logging_kwargs, parent_trace_id=sub_call_trace_id, + drop_invalid_edits=drop_invalid_edits, ) effective_messages: Final = polyfill_result.messages if polyfill_result is not None else messages @@ -841,6 +854,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: # bridge to it via ``run_async_function``. context_management: Final = kwargs.pop("context_management", None) additional_drop_params: Final[list[str] | None] = kwargs.get("additional_drop_params", None) + drop_invalid_edits: Final = litellm.drop_params is True or ( + "drop_params" in kwargs and kwargs["drop_params"] is True + ) # Deliberately do NOT auto-attach the proxy ``llm_router`` here: # ``run_async_function`` spawns a new event loop in a worker thread # to bridge to the async dispatcher, but the proxy router's httpx @@ -866,7 +882,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs) request_kwargs: Final = _REQUEST_KWARGS.validate_python(kwargs) polyfill_result = run_async_function( - _prepare_context_managed_request, + prepare_context_managed_request, model=model, messages=messages, tools=tools, @@ -878,6 +894,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: user_api_key_auth=user_api_key_auth, sub_call_logging_kwargs=proxy_sub_call_logging_kwargs(request_kwargs), parent_trace_id=proxy_sub_call_trace_id(request_kwargs), + drop_invalid_edits=drop_invalid_edits, ) effective_messages: Final = polyfill_result.messages if polyfill_result is not None else messages diff --git a/litellm/llms/anthropic/pass_through/context_management/dispatcher.py b/litellm/llms/anthropic/pass_through/context_management/dispatcher.py index 27b7412292e..3e1bab09d5a 100644 --- a/litellm/llms/anthropic/pass_through/context_management/dispatcher.py +++ b/litellm/llms/anthropic/pass_through/context_management/dispatcher.py @@ -85,6 +85,7 @@ async def apply_context_management( user_api_key_auth: "UserAPIKeyAuth | None" = None, sub_call_logging_kwargs: "ProxySubCallLoggingKwargs | None" = None, parent_trace_id: str | None = None, + count_only: bool = False, ) -> PolyfillResult: """Run edits in order; return a single ``PolyfillResult``. @@ -128,6 +129,7 @@ async def apply_context_management( user_api_key_auth=user_api_key_auth, sub_call_logging_kwargs=sub_call_logging_kwargs, parent_trace_id=parent_trace_id, + count_only=count_only, ) if editor_is_async else await asyncify(editor)( diff --git a/litellm/llms/anthropic/pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/pass_through/context_management/editors/compact.py index 09f49e9b663..5eb2863a241 100644 --- a/litellm/llms/anthropic/pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/pass_through/context_management/editors/compact.py @@ -642,6 +642,7 @@ async def apply_compact_20260112( user_api_key_auth: Optional["UserAPIKeyAuth"] = None, sub_call_logging_kwargs: ProxySubCallLoggingKwargs | None = None, parent_trace_id: str | None = None, + count_only: bool = False, ) -> PolyfillResult: """Apply ``compact_20260112``; return a ``PolyfillResult``. @@ -682,7 +683,7 @@ async def apply_compact_20260112( # Opt-in gate: no summary model configured → no-op (but still return the # Phase A-sliced/stripped messages so compaction blocks don't leak). - summary_model: Final = _read_summary_model_setting() + summary_model: Final = None if count_only else _read_summary_model_setting() if summary_model is None: applied["error"] = "summary_model_not_configured" # Slice-only forwarding: ``augmented_system`` already carries any prior diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index 0ed432cc394..d1f45751d3e 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -13,7 +13,7 @@ from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, ) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests azure_ai_anthropic_count_tokens_handler: Final = AzureAIAnthropicCountTokensHandler() @@ -54,6 +54,7 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): system: object | None = None, messages_format: TokenCountMessagesFormat = "anthropic", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: """ Count tokens using Azure AI Anthropic's CountTokens API. diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index 70ec6572675..5ed73e13188 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -16,7 +16,13 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk -from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import ( + CountContextManagement, + Message, + ProviderSpecificModelInfo, + TokenCountMessagesFormat, + TokenCountResponse, +) class BaseTokenCounter(ABC): @@ -32,6 +38,7 @@ class BaseTokenCounter(ABC): system: Any | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: pass diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 7ff2fddbb57..97f28b9fbc5 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -16,7 +16,13 @@ from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import ( + CountContextManagement, + LiteLLMPydanticObjectBase, + LlmProviders, + TokenCountMessagesFormat, + TokenCountResponse, +) RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api" MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api" @@ -109,6 +115,7 @@ class BedrockTokenCounter(BaseTokenCounter): system: object | None = None, messages_format: TokenCountMessagesFormat = "anthropic", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: if not messages: return None diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 3d04954dd41..4006c89954d 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -14,7 +14,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse GEMINI_IMAGE_ASPECT_RATIOS: Final[dict[str, float]] = { "1:1": 1 / 1, @@ -497,12 +497,14 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): system: object | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: import copy 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, @@ -511,7 +513,14 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): vertex_gemini_count_context, ) - request_view: Final = inference_request_view(messages or [], system, messages_format) + request_view: Final = await inference_request_view( + messages or [], + system, + messages_format, + count_context_edits(context_management, tools, (deployment or {}).get("litellm_params", {}), model_to_use), + ) + if isinstance(request_view, VertexGeminiCountContextError): + 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) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index 7c0697626ca..274d693baa7 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -29,7 +29,7 @@ from litellm.llms.openai.workload_identity import ( get_workload_identity_bearer_token, resolve_openai_workload_identity_config, ) -from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests openai_count_tokens_handler: Final = OpenAICountTokensHandler() @@ -95,6 +95,7 @@ class OpenAITokenCounter(BaseTokenCounter): system: object = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: """ Count tokens using OpenAI's Responses API /input_tokens endpoint. diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 581e2cb15f7..27ca4acf05d 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -23,7 +23,7 @@ from litellm.types.llms.vertex_ai import ( PartType, Schema, ) -from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages GEMINI_VIDEO_METADATA_KEYS: Final = MappingProxyType( @@ -1274,6 +1274,7 @@ class VertexAITokenCounter(BaseTokenCounter): system: object | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> TokenCountResponse | None: import copy @@ -1352,6 +1353,7 @@ class VertexAITokenCounter(BaseTokenCounter): else: from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContextError, + count_context_edits, count_context_error_response, gemini_count_contents, inference_request_view, @@ -1361,7 +1363,14 @@ class VertexAITokenCounter(BaseTokenCounter): ) from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter - request_view: Final = inference_request_view(messages or [], system, messages_format) + request_view: Final = await inference_request_view( + messages or [], + system, + messages_format, + count_context_edits(context_management, tools, count_tokens_params_request, model_to_use), + ) + if isinstance(request_view, VertexGeminiCountContextError): + 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) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index e094f4c5c0c..871532b7af5 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -10,11 +10,10 @@ from typing_extensions import assert_never 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 TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse if TYPE_CHECKING: from litellm.exceptions import APIError - from litellm.llms.anthropic.pass_through.context_management.result import PolyfillResult @dataclass(frozen=True, slots=True) @@ -26,6 +25,7 @@ class VertexGeminiCountContext: @dataclass(frozen=True, slots=True) class VertexGeminiCountContextError: error_message: str + status_code: int = 400 _SYSTEM_BLOCKS: Final = TypeAdapter(tuple[AnthropicMessagesTextParam, ...]) @@ -39,6 +39,11 @@ _VERTEX_TOOL_KEYS: Final = MappingProxyType( ) +_INVALID_SYSTEM: Final = VertexGeminiCountContextError( + error_message="Gemini system instructions only support text blocks." +) + + def _system_parts( system: object | None, ) -> tuple[Mapping[str, JsonValue], ...] | VertexGeminiCountContextError: @@ -49,7 +54,7 @@ def _system_parts( try: blocks: Final = _SYSTEM_BLOCKS.validate_python(system) except ValidationError: - return VertexGeminiCountContextError(error_message="Gemini system instructions only support text blocks.") + return _INVALID_SYSTEM return tuple({"text": block["text"]} for block in blocks) @@ -157,7 +162,7 @@ def count_context_error_response( tokenizer_type="gemini", error=True, error_message=error.error_message, - status_code=400, + status_code=error.status_code, ) @@ -179,36 +184,120 @@ class InferenceRequestView: system: object | None -def _client_compaction(messages: Sequence[Mapping[str, object]], system: object | None) -> "PolyfillResult | 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.pass_through.context_management.editors.compact import ( apply_client_compaction_block_history, ) - chat_messages: Final = [dict(message) for message in messages] - if system is None or isinstance(system, str): - return apply_client_compaction_block_history(messages=chat_messages, system=system) - try: - blocks: Final = _SYSTEM_BLOCKS.validate_python(system) - except ValidationError: - return None - return apply_client_compaction_block_history(messages=chat_messages, system=[dict(block) for block in blocks]) + 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) -def inference_request_view( - messages: Sequence[Mapping[str, object]], system: object | None, messages_format: TokenCountMessagesFormat -) -> InferenceRequestView: - """Apply the history rewrites /v1/messages runs before it converts a body for Gemini.""" +async def _anthropic_inference_view( + messages: Sequence[Mapping[str, object]], system: object | None, edits: CountContextEdits +) -> InferenceRequestView | VertexGeminiCountContextError: + from litellm.llms.anthropic.pass_through.adapters.handler import prepare_context_managed_request + from litellm.llms.anthropic.pass_through.context_management import AnthropicContextManagementError from litellm.llms.anthropic.pass_through.messages.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": - normalized: Final = normalize_native_messages(messages) - compacted: Final = _client_compaction(normalized, system) - if compacted is None: - return InferenceRequestView(messages=normalized, system=system) - return InferenceRequestView(messages=compacted.messages, system=compacted.system) + return await _anthropic_inference_view(messages, system, edits) case _: assert_never(messages_format) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index ee1ac79c820..a11268ad432 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -3,11 +3,13 @@ Unified /v1/messages endpoint - (Anthropic Spec) """ from collections.abc import Mapping +from types import MappingProxyType from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse from pydantic import JsonValue +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.anthropic_interface.exceptions import ( @@ -44,6 +46,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, with_litellm_call_id, ) +from litellm.proxy.litellm_pre_call_utils import should_auto_drop_params_for_agentic_cli router: Final = APIRouter() @@ -299,6 +302,26 @@ def _requested_output_format(data: Mapping[str, JsonValue]) -> JsonValue: return output_config.get("format") +_COUNTED_LITELLM_PARAMS: Final = ("additional_drop_params", "drop_params") + + +class _CountTokensBodyOptions(TypedDict): + output_format: ReadOnly[JsonValue] + context_management: ReadOnly[JsonValue] + request_litellm_params: ReadOnly[Mapping[str, JsonValue]] + + +def _count_tokens_body_options(data: Mapping[str, JsonValue], auto_drop_params: bool) -> _CountTokensBodyOptions: + sent_params: Final[Mapping[str, JsonValue]] = {key: data[key] for key in _COUNTED_LITELLM_PARAMS if key in data} + return { + "output_format": _requested_output_format(data), + "context_management": data.get("context_management"), + "request_litellm_params": MappingProxyType( + {"drop_params": True, **sent_params} if auto_drop_params else sent_params + ), + } + + @router.post( "/v1/messages/count_tokens", tags=["[beta] Anthropic Messages Token Counting"], @@ -328,7 +351,7 @@ async def count_tokens( Returns: {"input_tokens": } """ - from litellm.proxy.proxy_server import count_request_tokens + from litellm.proxy.proxy_server import count_request_tokens, proxy_config litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: @@ -359,7 +382,12 @@ async def count_tokens( request=token_request, call_endpoint=True, messages_format="anthropic", - output_format=_requested_output_format(data), + **_count_tokens_body_options( + data, + should_auto_drop_params_for_agentic_cli( + request.headers.get("user-agent", ""), "drop_params" in data, proxy_config + ), + ), ) except HTTPException as e: return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 675bf0830f9..3ad5caade1b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -945,7 +945,9 @@ def apply_missing_session_id_policy( ) -def should_auto_drop_params_for_agentic_cli(user_agent: str, data: dict, proxy_config: ProxyConfig) -> bool: +def should_auto_drop_params_for_agentic_cli( + user_agent: str, caller_set_drop_params: bool, proxy_config: ProxyConfig +) -> bool: """drop_params defaults to on for agentic CLIs so their client-specific params (e.g. Claude Code's thinking, Codex's service_tier) don't fail requests routed to providers that reject them. An explicit drop_params @@ -955,7 +957,7 @@ def should_auto_drop_params_for_agentic_cli(user_agent: str, data: dict, proxy_c if not (is_claude_code_user_agent(user_agent) or is_codex_user_agent(user_agent)): return False - if "drop_params" in data: + if caller_set_drop_params: return False config: Final = getattr(proxy_config, "config", None) litellm_settings: Final = config.get("litellm_settings") if isinstance(config, dict) else None @@ -2617,7 +2619,7 @@ async def add_litellm_data_to_request( user_agent = request.headers["user-agent"] data[_metadata_variable_name]["user_agent"] = user_agent - if should_auto_drop_params_for_agentic_cli(user_agent, data, proxy_config): + if should_auto_drop_params_for_agentic_cli(user_agent, "drop_params" in data, proxy_config): data["drop_params"] = True # Merge caller-supplied tags (x-litellm-tags header, data["tags"] root-level) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b56f0e68c4f..89fd3ae0f49 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -171,6 +171,7 @@ from litellm.router_utils.routing_groups import parse_routing_groups from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ( PRICING_OVERRIDES_KEY, + CountContextManagement, ModelResponse, ModelResponseStream, StreamingChoices, @@ -14456,6 +14457,7 @@ async def _try_provider_token_count( system: str | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: CountContextManagement | None = None, ) -> Optional["TokenCountResponse"]: """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting.""" if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider): @@ -14471,6 +14473,7 @@ async def _try_provider_token_count( system=system, messages_format=messages_format, output_format=output_format, + context_management=context_management, ) except httpx.HTTPStatusError as e: error_message: Final = getattr(e, "message", None) or str(e) @@ -14501,6 +14504,39 @@ def _drops_params(deployment: Mapping[str, object] | None) -> bool: return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True) +def _router_default_litellm_params(router: Router) -> Mapping[str, object]: + defaults: Final[Mapping[str, object] | None] = getattr(router, "default_litellm_params", None) + if defaults is None: + return {} + return {key: value for key, value in defaults.items() if value is not None} + + +def _count_context_management( + context_management: JsonValue, model: str, request_litellm_params: Mapping[str, JsonValue] +) -> CountContextManagement | None: + from litellm.llms.anthropic.pass_through.adapters.handler import context_management_explicitly_dropped + + if context_management is None or llm_router is None: + return None + inference_model: Final = resolve_inference_model( + 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} + for deployment in llm_router.get_routable_deployments( + inference_model if isinstance(inference_model, str) else model + ) + ) + if any(context_management_explicitly_dropped(params.get("additional_drop_params")) for params in sent_params): + return None + return CountContextManagement( + spec=context_management, + drop_invalid_edits=litellm.drop_params is True + or any(params.get("drop_params") is True for params in sent_params), + ) + + def _system_message(system: object) -> ChatCompletionSystemMessage | None: if not isinstance(system, (str, list)) or not system: return None @@ -14531,6 +14567,8 @@ async def count_request_tokens( call_endpoint: bool = False, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, + context_management: JsonValue = None, + request_litellm_params: Mapping[str, JsonValue] = MappingProxyType({}), ) -> TokenCountResponse: prompt: Final = request.prompt messages: Final = request.messages @@ -14602,6 +14640,7 @@ async def count_request_tokens( system=system, messages_format=messages_format, output_format=output_format, + context_management=_count_context_management(context_management, request.model, request_litellm_params), ) if result is not None: return result diff --git a/litellm/router.py b/litellm/router.py index 3c969610e83..4f2358f9008 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13727,6 +13727,15 @@ class Router: """ return [m for m in self.model_list if m["litellm_params"]["model"] == model] + def get_routable_deployments(self, model: str) -> tuple[DeploymentTypedDict, ...]: + try: + _, candidates = self._common_checks_available_deployment(model=model, request_kwargs={}) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # its signature leaves input and the result untyped + except (litellm.BadRequestError, ValueError): + return () + if isinstance(candidates, dict): + return (cast(DeploymentTypedDict, candidates),) # cast-ok: the router returns one deployment as a dict + return tuple(cast(list[DeploymentTypedDict], candidates)) # cast-ok: the router returns deployments as a list + def _try_early_resolve_deployments_for_model_not_in_names( self, model: str, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 7df8cc7452f..a97e35a9021 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2,6 +2,7 @@ import json import re import time from collections.abc import Collection, Mapping, Sequence +from dataclasses import dataclass from enum import Enum from types import MappingProxyType from typing import ( @@ -4477,6 +4478,12 @@ class LiteLLMLoggingBaseClass: TokenCountMessagesFormat: TypeAlias = Literal["openai", "anthropic"] +@dataclass(frozen=True, slots=True) +class CountContextManagement: + spec: JsonValue + drop_invalid_edits: bool + + class TokenCountResponse(LiteLLMPydanticObjectBase): total_tokens: int request_model: str diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 8de93256f6c..ec3fb941037 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -4,18 +4,22 @@ import threading from collections.abc import Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack +from dataclasses import dataclass +from itertools import chain +from pathlib import Path from types import MappingProxyType -from typing import Final +from typing import Final, Literal, TypeAlias, get_args from uuid import uuid4 import anthropic import httpx import pytest -from integration._support.client import Gateway, Scenario -from integration._support.process import graceful_stop_seconds +from integration._support.client import Gateway, Scenario, object_value, string_value +from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter +from typing_extensions import assert_never pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) @@ -302,6 +306,396 @@ def test_count_tokens_rejects_blocks_it_cannot_convert(gateway: Gateway, provide assert peer.drain() == () +_FILE_TEXT: Final = "line of a long file that the agent read earlier. " * 180 +_FILES: Final = 3 +_READ_FILE_TOOL: Final[JsonValue] = { + "name": "read_file", + "description": "Read a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, +} +_CLEAR_OLDER_TOOL_RESULTS: Final[JsonValue] = { + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, + } + ] +} +_COMPACT_UNDER_THE_MINIMUM_TRIGGER: Final[JsonValue] = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] +} + + +def _read_file_round(marker: str, index: int) -> tuple[JsonValue, ...]: + return ( + {"role": "user", "content": f"Read file {index}."}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}.txt"}} + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": f"toolu_{index}", + "content": f"{_FILE_TEXT}", + } + ], + }, + ) + + +def _agent_body(model: str, marker: str, options: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + rounds: Final = tuple(chain.from_iterable(_read_file_round(marker, index) for index in range(_FILES))) + return { + "model": model, + "tools": [_READ_FILE_TOOL], + "messages": [*rounds, {"role": "user", "content": "Summarize what you read."}], + "context_management": _CLEAR_OLDER_TOOL_RESULTS, + **options, + } + + +def _files_kept(contents: JsonValue, marker: str) -> tuple[bool, ...]: + sent: Final = json.dumps(contents) + return tuple(f"" in sent for index in range(_FILES)) + + +def _contents_sent_upstream( + gateway: Gateway, peer: _GeminiPeer, body: Mapping[str, JsonValue], headers: Mapping[str, str] | None = None +) -> tuple[JsonValue, JsonValue]: + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}, headers=headers) + assert generated.status_code == 200, generated.text + generate_requests: Final = tuple(request for request in peer.drain() if request.target.endswith(":generateContent")) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body, headers=headers) + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text + count_requests: Final = peer.drain() + assert len(generate_requests) == 1 + assert [request.target.endswith(":countTokens") for request in count_requests] == [True] + return _json_body(generate_requests[0])["contents"], _counted_contents(_json_body(count_requests[0])) + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_count_tokens_counts_the_history_messages_sends_after_context_edits(gateway: Gateway, provider: str) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _provider_deployment(gateway, scenario, provider, wire.url) + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream(gateway, peer, _agent_body(model, marker, {})) + assert _files_kept(generated, marker) == (False, False, True) + assert counted == generated + + +@pytest.mark.parametrize( + ("deployment_options", "request_options", "files_kept"), + ( + ({"additional_drop_params": ["context_management"]}, {}, (True, True, True)), + ({}, {"additional_drop_params": ["context_management"]}, (True, True, True)), + ({"additional_drop_params": ["context_management"]}, {"additional_drop_params": []}, (False, False, True)), + ), + ids=("deployment-drops", "request-drops", "request-overrides-the-deployment-drop"), +) +def test_count_tokens_follows_the_context_management_opt_out_messages_follows( + gateway: Gateway, + deployment_options: Mapping[str, JsonValue], + request_options: Mapping[str, JsonValue], + files_kept: tuple[bool, ...], +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url, **deployment_options) + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream(gateway, peer, _agent_body(model, marker, request_options)) + assert _files_kept(generated, marker) == files_kept + assert counted == generated + + +_RoutedCase: TypeAlias = Literal[ + "router-default-drops", + "fallback-deployment-drops", + "fallback-router-default-drops", + "fallback-request-drops", + "alias-to-a-litellm-model-drops", + "completion-model-drops", +] +_DROPS_CONTEXT_MANAGEMENT: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"additional_drop_params": ["context_management"]} +) + + +@dataclass(frozen=True, slots=True) +class _RoutedProxy: + settings: Mapping[str, JsonValue] + requested: str + request_options: Mapping[str, JsonValue] + + +def _scripted_gemini(model_name: str, model: str, api_base: str, options: Mapping[str, JsonValue]) -> JsonValue: + return { + "model_name": model_name, + "litellm_params": {"model": model, "api_base": api_base, "api_key": "scripted-gemini-key", **options}, + } + + +def _default_fallback_settings( + prefix: str, api_base: str, deployment_options: Mapping[str, JsonValue], router_defaults: Mapping[str, JsonValue] +) -> Mapping[str, JsonValue]: + return { + "model_list": [ + _scripted_gemini(f"{prefix}-fallback", f"gemini/{_MODEL}", api_base, deployment_options), + _scripted_gemini(f"{prefix}-wild/*", "gemini/*", api_base, {}), + ], + "router_settings": { + "default_fallbacks": [f"{prefix}-fallback"], + **({"default_litellm_params": dict(router_defaults)} if router_defaults else {}), + }, + } + + +def _routed_proxy(case: _RoutedCase, prefix: str, api_base: str) -> _RoutedProxy: + match case: + case "router-default-drops": + return _RoutedProxy( + { + "model_list": [_scripted_gemini(prefix, f"gemini/{_MODEL}", api_base, {})], + "router_settings": {"default_litellm_params": dict(_DROPS_CONTEXT_MANAGEMENT)}, + }, + prefix, + {}, + ) + case "fallback-deployment-drops": + return _RoutedProxy( + _default_fallback_settings(prefix, api_base, _DROPS_CONTEXT_MANAGEMENT, {}), f"{prefix}-unrouted", {} + ) + case "fallback-router-default-drops": + return _RoutedProxy( + _default_fallback_settings(prefix, api_base, {}, _DROPS_CONTEXT_MANAGEMENT), f"{prefix}-unrouted", {} + ) + case "fallback-request-drops": + return _RoutedProxy( + _default_fallback_settings(prefix, api_base, {}, {}), f"{prefix}-unrouted", _DROPS_CONTEXT_MANAGEMENT + ) + case "alias-to-a-litellm-model-drops": + return _RoutedProxy( + { + "model_list": [_scripted_gemini(prefix, f"gemini/{_MODEL}", api_base, _DROPS_CONTEXT_MANAGEMENT)], + "router_settings": {"model_group_alias": {f"{prefix}-alias": f"gemini/{_MODEL}"}}, + }, + f"{prefix}-alias", + {}, + ) + case "completion-model-drops": + return _RoutedProxy( + { + "model_list": [ + _scripted_gemini(f"{prefix}-keeps", f"gemini/{_MODEL}", api_base, {}), + _scripted_gemini(f"{prefix}-drops", f"gemini/{_MODEL}", api_base, _DROPS_CONTEXT_MANAGEMENT), + ], + "general_settings": {"completion_model": f"{prefix}-drops"}, + }, + f"{prefix}-keeps", + {}, + ) + case _: + assert_never(case) + + +@pytest.mark.parametrize("case", get_args(_RoutedCase)) +def test_count_tokens_follows_the_opt_out_of_the_deployment_proxy_settings_send_to( + gateway: Gateway, tmp_path: Path, case: _RoutedCase +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire: + routed: Final = _routed_proxy(case, f"routed-{uuid4().hex}", wire.url) + config: Final = tmp_path / "routed.yaml" + config.write_text(json.dumps(routed.settings), encoding="utf-8") + with owned_proxy_process(gateway, tmp_path, {}, config=config) as owned: + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream( + owned.gateway, peer, _agent_body(routed.requested, marker, routed.request_options) + ) + assert _files_kept(generated, marker) == (True, True, True) + assert counted == generated + + +def _named_deployment( + scenario: Scenario, + model_name: str, + litellm_params: Mapping[str, JsonValue], + model_info: Mapping[str, JsonValue], +) -> str: + created: Final = scenario.gateway.post( + "/model/new", + {"model_name": model_name, "litellm_params": dict(litellm_params), "model_info": dict(model_info)}, + ) + identity: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + return identity + + +def test_count_tokens_reads_the_opt_out_of_a_deployment_id_beside_a_wildcard_route(gateway: Gateway) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + prefix: Final = f"integration-{uuid4().hex}" + _named_deployment( + scenario, + f"{prefix}-*", + {"model": "gemini/*", "api_base": wire.url, "api_key": "scripted-gemini-key"}, + {}, + ) + deployment_id: Final = _named_deployment( + scenario, + prefix, + { + "model": f"gemini/{_MODEL}", + "api_base": wire.url, + "api_key": "scripted-gemini-key", + "additional_drop_params": ["context_management"], + }, + {"id": f"{prefix}-dropping"}, + ) + assert deployment_id == f"{prefix}-dropping" + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream(gateway, peer, _agent_body(deployment_id, marker, {})) + assert _files_kept(generated, marker) == (True, True, True) + assert counted == generated + + +@pytest.mark.parametrize( + ("requested", "wildcards", "files_kept"), + ( + ("dropping", (), (True, True, True)), + ("editing", (), (False, False, True)), + ("dropping", ("*",), (False, False, True)), + ), + ids=("deployment-drops", "sibling-deployment-drops", "wildcard-route-matches-first"), +) +def test_count_tokens_reads_the_opt_out_of_the_deployment_its_litellm_model_name_routes_to( + gateway: Gateway, requested: str, wildcards: tuple[str, ...], files_kept: tuple[bool, ...] +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + prefix: Final = f"gemini/{_MODEL}-{uuid4().hex}" + scenario.model( + model=f"{prefix}-dropping", + api_base=wire.url, + api_key="scripted-gemini-key", + additional_drop_params=["context_management"], + ) + scenario.model(model=f"{prefix}-editing", api_base=wire.url, api_key="scripted-gemini-key") + for wildcard in wildcards: + _named_deployment( + scenario, + f"{prefix}-{wildcard}", + {"model": f"{prefix}-{wildcard}", "api_base": wire.url, "api_key": "scripted-gemini-key"}, + {}, + ) + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream(gateway, peer, _agent_body(f"{prefix}-{requested}", marker, {})) + assert _files_kept(generated, marker) == files_kept + assert counted == generated + + +def _error_type(response: httpx.Response) -> JsonValue: + return object_value(_JSON_OBJECT.validate_json(response.content)["error"])["type"] + + +def _error_message(response: httpx.Response) -> JsonValue: + return object_value(_JSON_OBJECT.validate_json(response.content)["error"])["message"] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_count_tokens_rejects_a_context_edit_messages_rejects_without_an_upstream_call( + gateway: Gateway, provider: str +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _provider_deployment(gateway, scenario, provider, wire.url) + body: Final = _agent_body(model, uuid4().hex, {"context_management": _COMPACT_UNDER_THE_MINIMUM_TRIGGER}) + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 400, generated.text + assert counted.status_code == 400, counted.text + assert _error_type(counted) == _error_type(generated) + assert peer.drain() == () + + +def test_count_tokens_rejects_an_invalid_edit_from_a_coding_agent_that_turns_drop_params_off( + gateway: Gateway, +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url) + body: Final = _agent_body( + model, uuid4().hex, {"context_management": _COMPACT_UNDER_THE_MINIMUM_TRIGGER, "drop_params": False} + ) + claude_code: Final = MappingProxyType({"User-Agent": "claude-cli/2.1.0 (external, cli)"}) + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}, headers=claude_code) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body, headers=claude_code) + assert generated.status_code == 400, generated.text + assert counted.status_code == 400, counted.text + assert _error_type(counted) == _error_type(generated) + assert peer.drain() == () + + +@pytest.mark.parametrize( + ("deployment_options", "request_options", "user_agent"), + ( + ({"drop_params": True}, {}, "python-httpx"), + ({}, {"drop_params": True}, "python-httpx"), + ({}, {}, "claude-cli/2.1.0 (external, cli)"), + ({}, {}, "codex_cli_rs/0.50.0"), + ), + ids=("deployment-drop-params", "request-drop-params", "claude-code-client", "codex-client"), +) +def test_count_tokens_counts_the_unedited_body_messages_sends_when_drop_params_drops_an_invalid_edit( + gateway: Gateway, + deployment_options: Mapping[str, JsonValue], + request_options: Mapping[str, JsonValue], + user_agent: str, +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url, **deployment_options) + marker: Final = uuid4().hex + generated, counted = _contents_sent_upstream( + gateway, + peer, + _agent_body(model, marker, {"context_management": _COMPACT_UNDER_THE_MINIMUM_TRIGGER, **request_options}), + {"User-Agent": user_agent}, + ) + assert _files_kept(generated, marker) == (True, True, True) + assert counted == generated + + +def test_count_tokens_answers_for_the_unedited_body_when_a_context_edit_fails_inside_the_proxy( + gateway: Gateway, +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url) + unedited: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": _UNCONVERTIBLE_MESSAGES} + ) + edited: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": _UNCONVERTIBLE_MESSAGES, "context_management": _CLEAR_OLDER_TOOL_RESULTS}, + ) + assert unedited.status_code == 400, unedited.text + assert (edited.status_code, _error_type(edited), _error_message(edited)) == ( + unedited.status_code, + _error_type(unedited), + _error_message(unedited), + ) + assert peer.drain() == () + + def test_count_tokens_keeps_each_text_block_in_system_instruction(gateway: Gateway) -> None: peer: Final = _GeminiPeer() with wire_server(peer) as wire, gateway.scenario() as scenario: diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_context_management.py b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_context_management.py new file mode 100644 index 00000000000..f906530c8d9 --- /dev/null +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_context_management.py @@ -0,0 +1,131 @@ +from collections.abc import Mapping +from typing import Final +from unittest.mock import patch + +import pytest + +import litellm +from litellm.llms.anthropic.pass_through.adapters.handler import LiteLLMMessagesToCompletionTransformationHandler +from litellm.llms.anthropic.pass_through.context_management import AnthropicContextManagementError +from litellm.router import Router +from litellm.types.utils import Choices, Message, ModelResponse, Usage + +_TOOL_OUTPUT: Final = "line of the file under review " * 40 +_INVALID_COMPACT_EDIT: Final = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] +} + + +def _response() -> ModelResponse: + return ModelResponse( + choices=[Choices(index=0, finish_reason="stop", message=Message(content="done"))], + usage=Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11), + ) + + +def _messages() -> list[dict[str, object]]: + return [ + {"role": "user", "content": "review a.py"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "read_file", "input": {"path": "a.py"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": _TOOL_OUTPUT}]}, + ] + + +class _Backend: + def __init__(self) -> None: + self.calls: list[Mapping[str, object]] = [] + + async def __call__(self, **kwargs: object) -> ModelResponse: + self.calls.append(kwargs) + return _response() + + +class _SyncBackend: + def __init__(self) -> None: + self.calls: list[Mapping[str, object]] = [] + + def __call__(self, **kwargs: object) -> ModelResponse: + self.calls.append(kwargs) + return _response() + + +def _router(deployment_params: Mapping[str, object], router_defaults: Mapping[str, object]) -> Router: + return Router( + model_list=[ + { + "model_name": "gemini-agent", + "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "test-key", **deployment_params}, + } + ], + default_litellm_params=dict(router_defaults), + ) + + +async def _send(router: Router, request_options: Mapping[str, object], backend: _Backend) -> object: + with patch("litellm.acompletion", new=backend): + return await router.aanthropic_messages( + model="gemini-agent", + max_tokens=16, + messages=_messages(), + tools=[{"name": "read_file", "input_schema": {"type": "object", "properties": {}}}], + context_management=_INVALID_COMPACT_EDIT, + **request_options, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_params", "router_defaults", "request_options", "litellm_drop_params"), + ( + ({"drop_params": True}, {}, {}, False), + ({}, {"drop_params": True}, {}, False), + ({}, {}, {"drop_params": True}, False), + ({}, {}, {}, True), + ), + ids=("deployment-drop-params", "router-default-drop-params", "request-drop-params", "litellm-drop-params"), +) +async def test_messages_sends_the_unedited_body_when_drop_params_drops_an_invalid_edit( + deployment_params: Mapping[str, object], + router_defaults: Mapping[str, object], + request_options: Mapping[str, object], + litellm_drop_params: bool, +) -> None: + backend: Final = _Backend() + with patch.object(litellm, "drop_params", litellm_drop_params): + response: Final = await _send(_router(deployment_params, router_defaults), request_options, backend) + + assert isinstance(response, dict) + assert response["content"] == [{"type": "text", "text": "done"}] + assert len(backend.calls) == 1 + assert _TOOL_OUTPUT in str(backend.calls[0]["messages"]) + + +@pytest.mark.asyncio +async def test_messages_rejects_an_invalid_edit_without_drop_params() -> None: + backend: Final = _Backend() + with patch.object(litellm, "drop_params", False), pytest.raises(AnthropicContextManagementError) as error: + await _send(_router({}, {}), {}, backend) + + assert error.value.status_code == 400 + assert backend.calls == [] + + +def test_sync_messages_sends_the_unedited_body_when_drop_params_drops_an_invalid_edit() -> None: + backend: Final = _SyncBackend() + with patch.object(litellm, "drop_params", False), patch("litellm.completion", new=backend): + response: Final = LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( + max_tokens=16, + messages=_messages(), + model="gemini/gemini-2.5-flash", + api_key="test-key", + context_management=_INVALID_COMPACT_EDIT, + drop_params=True, + ) + + assert isinstance(response, dict) + assert response["content"] == [{"type": "text", "text": "done"}] + assert len(backend.calls) == 1 + assert _TOOL_OUTPUT in str(backend.calls[0]["messages"]) diff --git a/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py b/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py index 0a5fbcdf63e..5e55c0ade32 100644 --- a/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py +++ b/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py @@ -2544,7 +2544,7 @@ async def test_prepare_context_managed_request_forwards_proxy_litellm_metadata() the summary subcall lands on the router with no parent attribution, and those tokens go unbilled to the caller's key/team.""" from litellm.llms.anthropic.pass_through.adapters.handler import ( - _prepare_context_managed_request, + prepare_context_managed_request, ) captured_summary_metadata: Dict[str, Any] = {} @@ -2561,7 +2561,7 @@ async def test_prepare_context_managed_request_forwards_proxy_litellm_metadata() ), patch("litellm.token_counter", return_value=200_000), ): - result = await _prepare_context_managed_request( + result = await prepare_context_managed_request( model=MODEL, messages=_simple_messages(), tools=None, diff --git a/tests/unit/llms/vertex_ai/count_tokens/test_context.py b/tests/unit/llms/vertex_ai/count_tokens/test_context.py index 06dacedde67..405b7234c81 100644 --- a/tests/unit/llms/vertex_ai/count_tokens/test_context.py +++ b/tests/unit/llms/vertex_ai/count_tokens/test_context.py @@ -1,5 +1,6 @@ import asyncio import base64 +import logging import time import uuid from collections.abc import Generator @@ -18,10 +19,13 @@ from litellm.llms.anthropic.pass_through.adapters.transformation import ( truncate_tool_name, ) from litellm.llms.vertex_ai.count_tokens.context import ( + CountContextEdits, + InferenceRequestView, VertexGeminiCountContext, VertexGeminiCountContextError, gemini_count_contents, inference_request_view, + unconvertible_messages_error, vertex_gemini_count_context, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig @@ -31,6 +35,7 @@ _SCHEMA: Final[JsonValue] = { "properties": {}, "additionalProperties": False, } +_NO_EDITS: Final = CountContextEdits(spec=None, model="gemini/gemini-2.5-flash", tools=None) _FUNCTION_DECLARATION: Final = { "functionDeclarations": [ { @@ -383,8 +388,8 @@ async def test_image_that_cannot_be_fetched_returns_an_error_value( @pytest.mark.parametrize( "system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system") ) -def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None: - view: Final = inference_request_view( +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."}]}, @@ -392,16 +397,18 @@ def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(sys ), 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) -def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None: +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 = inference_request_view( + view: Final = await inference_request_view( ( {"role": "user", "content": "What is new?"}, { @@ -426,18 +433,21 @@ def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> Non ), 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"] -def test_openai_view_is_the_request_unchanged() -> None: +async def test_openai_view_is_the_request_unchanged() -> None: messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},) - view: Final = inference_request_view(messages, "Be brief.", "openai") + 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." @@ -470,7 +480,7 @@ async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> Non @pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini")) -def test_anthropic_replayed_web_search_result_is_counted_as_its_text( +async def test_anthropic_replayed_web_search_result_is_counted_as_its_text( custom_llm_provider: Literal["vertex_ai", "gemini"], ) -> None: snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." @@ -498,9 +508,163 @@ def test_anthropic_replayed_web_search_result_is_counted_as_its_text( {"role": "user", "content": "Thanks."}, ) - view: Final = inference_request_view(messages, None, "anthropic") - contents: Final = asyncio.run( - gemini_count_contents(view.messages, "anthropic", "gemini-2.5-flash", custom_llm_provider, {}) + view: Final = await inference_request_view(messages, None, "anthropic", _NO_EDITS) + assert isinstance(view, InferenceRequestView) + contents: Final = await gemini_count_contents( + view.messages, "anthropic", "gemini-2.5-flash", custom_llm_provider, {} ) 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.pass_through.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_before_compact_fails() -> 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), + ) + + 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: + turns: Final = ( + *_agent_turns(1)[:-1], + { + "role": "user", + "content": [ + { + "type": "search_result", + "source": "https://docs.litellm.ai", + "title": "Docs", + "content": [{"type": "text", "text": "LiteLLM docs"}], + } + ], + }, + ) + + view: Final = await inference_request_view(turns, None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) + + assert isinstance(view, InferenceRequestView) + blocks_error: Final = unconvertible_messages_error(view.messages, "anthropic") + assert blocks_error is not None + assert blocks_error.status_code == 400 + assert "search_result" in blocks_error.error_message diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index d244fe6bd90..97c0a657ca4 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -28,7 +28,7 @@ from litellm.llms.vertex_ai.common_utils import ( vertex_request_labels_from_litellm_params, ) from litellm.llms.vertex_ai.count_tokens import handler as vertex_count_tokens_handler -from litellm.types.utils import TokenCountMessagesFormat +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat @pytest.mark.asyncio @@ -62,7 +62,9 @@ class _VertexCountTokensPeer: return httpx.Response(200, json={"totalTokens": 37}, request=request) return httpx.Response( self.status_code, - json={"error": {"code": self.status_code, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}, + json={ + "error": {"code": self.status_code, "message": "Invalid function name", "status": "INVALID_ARGUMENT"} + }, request=request, ) @@ -133,6 +135,60 @@ async def test_vertex_gemini_count_rejects_an_image_url_it_refuses_to_fetch_with assert peer.requests == () +_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180 + + +def _tool_round(index: int) -> tuple[dict[str, object], ...]: + return ( + {"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}], + }, + ) + + +@pytest.mark.asyncio +async def test_vertex_gemini_count_applies_the_context_edits_inference_applies() -> None: + with _vertex_count_tokens_peer(200) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[*_tool_round(0), *_tool_round(1), *_tool_round(2), {"role": "user", "content": "Summarize."}], + contents=None, + deployment={ + "litellm_params": {**_VERTEX_DEPLOYMENT["litellm_params"], "model": "vertex_ai/gemini-2.5-flash"} + }, + request_model="vertex_ai/gemini-2.5-flash", + tools=[ + {"name": "read_file", "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}} + ], + messages_format="anthropic", + context_management=CountContextManagement( + spec={ + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, + } + ] + }, + drop_invalid_edits=False, + ), + ) + + assert result is not None + assert result.total_tokens == 37 + assert len(peer.requests) == 1 + assert peer.requests[0].content.decode().count(_TOOL_OUTPUT) == 1 + + @pytest.mark.parametrize("location", ["us", "eu", "us-central1", "europe-west1", "global"]) def test_validate_vertex_location_accepts_valid(location): assert validate_vertex_location(location) == location diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 7d70e058396..f7d1970fdae 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -3,6 +3,7 @@ import json +from collections.abc import Callable, Mapping from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -153,6 +154,7 @@ async def test_anthropic_messages_count_tokens_endpoint(): # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-3-sonnet-20240229", "messages": [{"role": "user", "content": "Hello Claude!"}], @@ -172,7 +174,14 @@ async def test_anthropic_messages_count_tokens_endpoint(): anthropic_endpoints.read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=False, messages_format="openai", output_format=None): + async def mock_token_counter( + request, + call_endpoint=False, + messages_format="openai", + output_format=None, + context_management=None, + request_litellm_params=None, + ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "claude-3-sonnet-20240229" assert request.messages == [{"role": "user", "content": "Hello Claude!"}] @@ -225,6 +234,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "gpt-4", "messages": [{"role": "user", "content": "Hello GPT!"}], @@ -244,7 +254,14 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): anthropic_endpoints.read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=True, messages_format="openai", output_format=None): + async def mock_token_counter( + request, + call_endpoint=True, + messages_format="openai", + output_format=None, + context_management=None, + request_litellm_params=None, + ): assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "gpt-4" assert request.messages == [{"role": "user", "content": "Hello GPT!"}] @@ -1055,7 +1072,14 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with Bedrock-style error - async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): + async def mock_token_counter_error( + request, + call_endpoint=False, + messages_format="openai", + output_format=None, + context_management=None, + request_litellm_params=None, + ): raise ProxyException( message='{"detail":{"message":"Input is too long for requested model."}}', type="token_counting_error", @@ -1105,7 +1129,14 @@ async def test_anthropic_endpoint_403_permission_error_format(): original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 403 error - async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): + async def mock_token_counter_error( + request, + call_endpoint=False, + messages_format="openai", + output_format=None, + context_management=None, + request_litellm_params=None, + ): raise ProxyException( message='{"Message":"Bearer Token has expired"}', type="token_counting_error", @@ -1155,7 +1186,14 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 429 error - async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): + async def mock_token_counter_error( + request, + call_endpoint=False, + messages_format="openai", + output_format=None, + context_management=None, + request_litellm_params=None, + ): raise ProxyException( message="Rate limit exceeded", type="token_counting_error", @@ -1179,6 +1217,310 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): proxy_server.count_request_tokens = original_token_counter +_GEMINI_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180 + + +class _GeminiCountTokensPeer: + def __init__(self) -> None: + self.request_bodies: tuple[bytes, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.request_bodies = (*self.request_bodies, request.content) + return httpx.Response(200, json={"totalTokens": 41}, request=request) + + +def _agent_turns_with_tool_results(rounds: int) -> list[dict[str, JsonValue]]: + return [ + *( + turn + for index in range(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": _GEMINI_TOOL_OUTPUT} + ], + }, + ) + ), # comprehension-ok: flattens each round into its three turns + {"role": "user", "content": "Summarize what you read."}, + ] + + +_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}, + } + ] +} +_INVALID_COMPACT_EDIT: Final[dict[str, JsonValue]] = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] +} + + +def _gemini_deployment(model_name: str, litellm_params: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return { + "model_name": model_name, + "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "test-key", **litellm_params}, + } + + +async def _count_through_the_messages_endpoint( + router: Router, request_options: Mapping[str, JsonValue] +) -> tuple[object, tuple[bytes, ...]]: + import litellm.proxy.anthropic_endpoints.endpoints as anthropic_endpoints + import litellm.proxy.proxy_server as proxy_server + from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler + + peer: Final = _GeminiCountTokensPeer() + request_body: Final[dict[str, JsonValue]] = { + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": _CLEAR_TOOL_USES, + **request_options, + } + 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=request_body)), + patch.object( + gemini_count_tokens_handler, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ), + ): + response: Final = await anthropic_count_tokens(mock_request, MagicMock()) + return response, peer.request_bodies + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ( + "requested_model", + "additional_drop_params", + "router_additional_drop_params", + "request_options", + "expected_tool_outputs_sent", + ), + ( + ("gemini-agent", None, None, {}, 1), + ("gemini-agent", ["context_management"], None, {}, 3), + ("gemini-agent", None, None, {"additional_drop_params": ["context_management"]}, 3), + ("gemini-agent", ["context_management"], None, {"additional_drop_params": []}, 1), + ("gemini-agent", None, ["context_management"], {}, 3), + ("gemini-agent", None, ["context_management"], {"additional_drop_params": []}, 1), + ("gemini/gemini-2.5-flash", None, None, {}, 1), + ("gemini/gemini-2.5-flash", ["context_management"], None, {}, 3), + ("gemini/gemini-2.5-flash", None, ["context_management"], {}, 3), + ("gemini/gemini-2.5-flash", None, None, {"additional_drop_params": ["context_management"]}, 3), + ("gemini/gemini-2.5-flash", ["context_management"], None, {"additional_drop_params": []}, 1), + ), + ids=( + "edits-applied", + "deployment-drops-context-management", + "request-drops-context-management", + "request-overrides-the-deployment-drop", + "router-default-drops-context-management", + "request-overrides-the-router-default-drop", + "litellm-model-name-edits-applied", + "litellm-model-name-deployment-drops-context-management", + "litellm-model-name-router-default-drops-context-management", + "litellm-model-name-request-drops-context-management", + "litellm-model-name-request-overrides-the-deployment-drop", + ), +) +async def test_gemini_messages_count_tokens_applies_the_context_edits_inference_applies( + requested_model: str, + additional_drop_params: list[str] | None, + router_additional_drop_params: list[str] | None, + request_options: dict[str, JsonValue], + expected_tool_outputs_sent: int, +) -> None: + router: Final = Router( + model_list=[_gemini_deployment("gemini-agent", {"additional_drop_params": additional_drop_params})], + default_litellm_params={"additional_drop_params": router_additional_drop_params}, + ) + response, sent_bodies = await _count_through_the_messages_endpoint( + router, {**request_options, "model": requested_model} + ) + + assert response == {"input_tokens": 41} + assert len(sent_bodies) == 1 + assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == expected_tool_outputs_sent + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("requested_model", "model_list"), + ( + ( + "gemini-agent", + [ + { + "model_name": "gemini-agent", + "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "test-key", "order": 1}, + }, + { + "model_name": "gemini-agent", + "litellm_params": { + "model": "gemini/gemini-2.5-flash", + "api_key": "test-key", + "order": 2, + "additional_drop_params": ["context_management"], + }, + }, + ], + ), + ( + "gemini-agent-deployment-id", + [ + { + "model_name": "gemini-agent", + "litellm_params": { + "model": "gemini/gemini-2.5-flash", + "api_key": "test-key", + "additional_drop_params": ["context_management"], + }, + "model_info": {"id": "gemini-agent-deployment-id"}, + }, + ], + ), + ( + "gemini-agent-deployment-id", + [ + {"model_name": "*", "litellm_params": {"model": "gemini/*", "api_key": "test-key"}}, + { + "model_name": "gemini-agent", + "litellm_params": { + "model": "gemini/gemini-2.5-flash", + "api_key": "test-key", + "additional_drop_params": ["context_management"], + }, + "model_info": {"id": "gemini-agent-deployment-id"}, + }, + ], + ), + ), + ids=("group-member-drops", "deployment-id-drops", "deployment-id-drops-beside-a-wildcard"), +) +async def test_gemini_messages_count_tokens_keeps_tool_results_when_a_routable_deployment_drops_context_management( + requested_model: str, model_list: list[dict[str, JsonValue]] +) -> None: + response, sent_bodies = await _count_through_the_messages_endpoint( + Router(model_list=model_list), {"model": requested_model} + ) + + assert response == {"input_tokens": 41} + assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("requested_model", "make_router", "request_options", "expected_tool_outputs_sent"), + ( + ( + "unrouted-model", + lambda: Router(model_list=[_gemini_deployment("gemini-agent", {})], default_fallbacks=["gemini-agent"]), + {}, + 1, + ), + ( + "unrouted-model", + lambda: Router( + model_list=[_gemini_deployment("gemini-agent", {"additional_drop_params": ["context_management"]})], + default_fallbacks=["gemini-agent"], + ), + {}, + 3, + ), + ( + "unrouted-model", + lambda: Router( + model_list=[_gemini_deployment("gemini-agent", {})], + default_fallbacks=["gemini-agent"], + default_litellm_params={"additional_drop_params": ["context_management"]}, + ), + {}, + 3, + ), + ( + "unrouted-model", + lambda: Router(model_list=[_gemini_deployment("gemini-agent", {})], default_fallbacks=["gemini-agent"]), + {"additional_drop_params": ["context_management"]}, + 3, + ), + ( + "flash-alias", + lambda: Router( + model_list=[_gemini_deployment("gemini-agent", {"additional_drop_params": ["context_management"]})], + model_group_alias={"flash-alias": "gemini/gemini-2.5-flash"}, + ), + {}, + 3, + ), + ), + ids=( + "default-fallback-edits-applied", + "default-fallback-deployment-drops-context-management", + "default-fallback-router-default-drops-context-management", + "default-fallback-request-drops-context-management", + "alias-to-a-litellm-model-deployment-drops-context-management", + ), +) +async def test_gemini_messages_count_tokens_reads_the_opt_out_of_the_deployment_the_router_falls_back_to( + requested_model: str, + make_router: Callable[[], Router], + request_options: dict[str, JsonValue], + expected_tool_outputs_sent: int, +) -> None: + response, sent_bodies = await _count_through_the_messages_endpoint( + make_router(), {**request_options, "model": requested_model} + ) + + assert response == {"input_tokens": 41} + assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == expected_tool_outputs_sent + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("requested_model", "completion_model", "expected_tool_outputs_sent"), + (("gemini-agent", "gemini-dropping", 3), ("gemini-dropping", "gemini-agent", 1)), + ids=("completion-model-drops-context-management", "completion-model-keeps-context-management"), +) +async def test_gemini_messages_count_tokens_reads_the_opt_out_of_the_completion_model_inference_sends_to( + requested_model: str, completion_model: str, expected_tool_outputs_sent: int +) -> None: + import litellm.proxy.proxy_server as proxy_server + + router: Final = Router( + model_list=[ + _gemini_deployment("gemini-agent", {}), + _gemini_deployment("gemini-dropping", {"additional_drop_params": ["context_management"]}), + ] + ) + with patch.object(proxy_server, "general_settings", {"completion_model": completion_model}): + response, sent_bodies = await _count_through_the_messages_endpoint(router, {"model": requested_model}) + + assert response == {"input_tokens": 41} + assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == expected_tool_outputs_sent + + class _GeminiPart(BaseModel): model_config = ConfigDict(frozen=True) text: str = "" @@ -1224,6 +1566,48 @@ def _gemini_router(drop_params: bool = False) -> Router: ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_params", "router_defaults", "request_options", "litellm_drop_params"), + ( + ({"drop_params": True}, {}, {}, False), + ({}, {"drop_params": True}, {}, False), + ({}, {}, {"drop_params": True}, False), + ({}, {}, {}, True), + ), + ids=("deployment-drop-params", "router-default-drop-params", "request-drop-params", "litellm-drop-params"), +) +async def test_gemini_messages_count_tokens_counts_the_unedited_body_when_drop_params_drops_an_invalid_edit( + deployment_params: dict[str, JsonValue], + router_defaults: dict[str, JsonValue], + request_options: dict[str, JsonValue], + litellm_drop_params: bool, +) -> None: + router: Final = Router( + model_list=[_gemini_deployment("gemini-agent", deployment_params)], default_litellm_params=router_defaults + ) + with patch.object(litellm, "drop_params", litellm_drop_params): + response, sent_bodies = await _count_through_the_messages_endpoint( + router, {**request_options, "model": "gemini-agent", "context_management": _INVALID_COMPACT_EDIT} + ) + + assert response == {"input_tokens": 41} + assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == 3 + + +@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): + response, sent_bodies = await _count_through_the_messages_endpoint( + Router(model_list=[_gemini_deployment("gemini-agent", {})]), + {"model": "gemini-agent", "context_management": _INVALID_COMPACT_EDIT}, + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert sent_bodies == () + + @pytest.mark.asyncio @pytest.mark.parametrize( "text", ["hello there", "The quick brown fox jumps over the lazy dog. " * 50], ids=["short", "long"] diff --git a/tests/unit/test_router/test_router_get_deployments.py b/tests/unit/test_router/test_router_get_deployments.py index 8b0c34f863c..07486762a64 100644 --- a/tests/unit/test_router/test_router_get_deployments.py +++ b/tests/unit/test_router/test_router_get_deployments.py @@ -5,6 +5,7 @@ import os import random import traceback from collections import defaultdict +from collections.abc import Callable from typing import Final, Literal import pytest @@ -110,9 +111,7 @@ def test_get_available_deployment_for_pass_through(): ) # Test that only pass-through deployment is returned - selected_model = router.get_available_deployment_for_pass_through( - "gpt-3.5-turbo" - ) + selected_model = router.get_available_deployment_for_pass_through("gpt-3.5-turbo") assert selected_model["litellm_params"]["model"] == "gpt-3.5-turbo" assert selected_model["litellm_params"]["use_in_pass_through"] is True @@ -259,9 +258,7 @@ def test_filter_pass_through_deployments(): healthy_deployments = router.get_model_list() # Filter pass-through deployments - pass_through_deployments = router._filter_pass_through_deployments( - healthy_deployments - ) + pass_through_deployments = router._filter_pass_through_deployments(healthy_deployments) # Should only have 2 deployments with use_in_pass_through=True assert len(pass_through_deployments) == 2 @@ -280,9 +277,7 @@ async def _weighted_async_deployment_ids(router: Router) -> tuple[str, ...]: random.seed(2025) deployments: Final = tuple( [ - await router.async_get_available_deployment( - model="shared", messages=None, request_kwargs={} - ) + await router.async_get_available_deployment(model="shared", messages=None, request_kwargs={}) for _ in range(1000) ] ) @@ -326,10 +321,7 @@ def test_weighted_selection_router( selected_ids: Final = ( asyncio.run(_weighted_async_deployment_ids(router)) if async_selection - else tuple( - router.get_available_deployment("shared")["model_info"]["id"] - for _ in range(1000) - ) + else tuple(router.get_available_deployment("shared")["model_info"]["id"] for _ in range(1000)) ) assert selected_ids.count("high") / len(selected_ids) > 0.89 @@ -355,10 +347,7 @@ def test_weighted_selection_router_no_rpm_set() -> None: }, ] ) - selected_ids: Final = tuple( - router.get_available_deployment("shared")["model_info"]["id"] - for _ in range(100) - ) + selected_ids: Final = tuple(router.get_available_deployment("shared")["model_info"]["id"] for _ in range(100)) assert set(selected_ids) == {"low", "high"} @@ -378,3 +367,61 @@ def test_model_group_aliases_select_the_target_group() -> None: assert deployment["model_name"] == "target" assert deployment["model_info"]["id"] == "target-deployment" + + +_ROUTABLE_MODEL_LIST: Final = [ + { + "model_name": "shared", + "litellm_params": {"model": "openai/shared-low", "api_key": "test-key"}, + "model_info": {"id": "low"}, + }, + { + "model_name": "shared", + "litellm_params": {"model": "openai/shared-high", "api_key": "test-key"}, + "model_info": {"id": "high"}, + }, + { + "model_name": "unrelated", + "litellm_params": {"model": "openai/unrelated", "api_key": "test-key"}, + "model_info": {"id": "unrelated"}, + }, +] + + +@pytest.mark.parametrize( + ("requested_model", "make_router", "expected_ids"), + ( + ("shared", lambda: Router(model_list=_ROUTABLE_MODEL_LIST), {"low", "high"}), + ("high", lambda: Router(model_list=_ROUTABLE_MODEL_LIST), {"high"}), + ("openai/unrelated", lambda: Router(model_list=_ROUTABLE_MODEL_LIST), {"unrelated"}), + ( + "unknown", + lambda: Router(model_list=_ROUTABLE_MODEL_LIST, default_fallbacks=["shared"]), + {"low", "high"}, + ), + ( + "alias", + lambda: Router(model_list=_ROUTABLE_MODEL_LIST, model_group_alias={"alias": "openai/unrelated"}), + {"unrelated"}, + ), + ("unknown", lambda: Router(model_list=_ROUTABLE_MODEL_LIST), set()), + ), + ids=( + "model-group", + "deployment-id", + "litellm-model-name", + "default-fallback", + "alias-to-a-litellm-model-name", + "unroutable", + ), +) +def test_get_routable_deployments_returns_every_deployment_routing_can_pick( + requested_model: str, make_router: Callable[[], Router], expected_ids: set[str] +) -> None: + router: Final = make_router() + + routable_ids: Final = { + deployment["model_info"]["id"] for deployment in router.get_routable_deployments(requested_model) + } + + assert routable_ids == expected_ids