diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index a45d6c1184f..d78de0037a0 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -10,7 +10,7 @@ from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests anthropic_count_tokens_handler: Final = AnthropicCountTokensHandler() @@ -34,6 +34,8 @@ class AnthropicTokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, JsonValue]] | None = None, system: JsonValue = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using Anthropic's CountTokens API. diff --git a/litellm/llms/anthropic/pass_through/messages/utils.py b/litellm/llms/anthropic/pass_through/messages/utils.py index ae5ab68242b..d8b96c1b7a0 100644 --- a/litellm/llms/anthropic/pass_through/messages/utils.py +++ b/litellm/llms/anthropic/pass_through/messages/utils.py @@ -129,6 +129,13 @@ def anthropic_system_to_openai_message(system: object) -> ChatCompletionSystemMe return ChatCompletionSystemMessage(role="system", content=system) +def normalize_native_messages(messages: Sequence[object]) -> list[dict[str, object]]: + """Apply the block rewrites every /v1/messages request gets before it reaches a provider.""" + return flatten_unencrypted_web_search_results_in_anthropic_messages( + sanitize_tool_use_ids_in_anthropic_messages(strip_empty_content_blocks_from_anthropic_messages(messages)) + ) + + def prepare_native_messages( messages: list[dict[str, JsonValue]], system: str | list[dict[str, JsonValue]] | None, @@ -140,13 +147,7 @@ def prepare_native_messages( api_base: str | None = None, presanitized: bool = False, ) -> tuple[list[dict[str, JsonValue]], str | list[dict[str, JsonValue]] | None]: - normalized: Final = ( - messages - if presanitized - else flatten_unencrypted_web_search_results_in_anthropic_messages( - sanitize_tool_use_ids_in_anthropic_messages(strip_empty_content_blocks_from_anthropic_messages(messages)) - ) - ) + normalized: Final = messages if presanitized else normalize_native_messages(messages) return cast( # cast-ok: legacy normalizers and injection preserve the JSON message and system shapes tuple[list[dict[str, JsonValue]], str | list[dict[str, JsonValue]] | None], AnthropicCacheControlHook.maybe_inject_cache_control( diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py index 156c8d89d81..97f2db36a85 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py @@ -579,6 +579,14 @@ class LiteLLMAnthropicToResponsesAPIAdapter: result["summary"] = "detailed" return result + @classmethod + def system_blocks_to_instructions(cls, system: object) -> str: + return "\n".join( + text + for block in cls._content_blocks(system) + if block.get("type") == "text" and isinstance(text := block.get("text"), str) and text + ) + def translate_request( self, anthropic_request: AnthropicMessagesRequest, @@ -628,9 +636,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if isinstance(system, str): responses_kwargs["instructions"] = system elif isinstance(system, list): - responses_kwargs["instructions"] = "\n".join( - filter(None, (b.get("text", "") for b in system if isinstance(b, dict) and b.get("type") == "text")) - ) + responses_kwargs["instructions"] = self.system_blocks_to_instructions(system) # max_tokens -> max_output_tokens max_tokens: Final = anthropic_request.get("max_tokens") 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 1d6d59c0592..2707d24ee07 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -12,7 +12,7 @@ from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, ) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests azure_ai_anthropic_count_tokens_handler: Final = AzureAIAnthropicCountTokensHandler() @@ -36,6 +36,8 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, JsonValue]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using Azure AI Anthropic's CountTokens API. diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index fee12e34fa3..70ec6572675 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -16,7 +16,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk -from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountResponse +from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountMessagesFormat, TokenCountResponse class BaseTokenCounter(ABC): @@ -30,6 +30,8 @@ class BaseTokenCounter(ABC): request_model: str = "", tools: list[dict[str, Any]] | None = None, system: Any | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: pass diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 3084b0fde07..7ff2fddbb57 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -16,7 +16,7 @@ from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountResponse +from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountMessagesFormat, TokenCountResponse RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api" MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api" @@ -107,6 +107,8 @@ class BedrockTokenCounter(BaseTokenCounter): request_model: str = "", tools: Sequence[Mapping[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "anthropic", + output_format: JsonValue = None, ) -> TokenCountResponse | None: if not messages: return None diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index b18c322b3a4..57ee0f33879 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence from typing import Any, Final import httpx -from pydantic import TypeAdapter +from pydantic import JsonValue, TypeAdapter import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -14,7 +14,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import TokenCountResponse +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse GEMINI_IMAGE_ASPECT_RATIOS: Final[dict[str, float]] = { "1:1": 1 / 1, @@ -495,6 +495,8 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: import copy diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 02a05b060b8..6f911e9c40f 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -5,7 +5,7 @@ Uses httpx for HTTP requests to OpenAI's /v1/responses/input_tokens endpoint. """ import json -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Final import httpx @@ -31,8 +31,9 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): api_key: str, api_base: str | None = None, timeout: float | httpx.Timeout | None = None, - tools: list[dict[str, object]] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, instructions: str | None = None, + text: Mapping[str, object] | None = None, ) -> dict[str, object]: """ Handle a token counting request to OpenAI's Responses API. @@ -53,6 +54,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): input=input, tools=tools, instructions=instructions, + text=text, ) endpoint_url: Final = self.get_openai_count_tokens_endpoint(api_base) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index ab3ddbd179d..b97e1329675 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -3,9 +3,16 @@ OpenAI Token Counter implementation using the Responses API /input_tokens endpoi """ import os +from collections.abc import Mapping, Sequence from typing import Any, Final +from pydantic import JsonValue +from typing_extensions import assert_never + from litellm._logging import verbose_logger +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.responses.count_tokens.handler import ( @@ -13,8 +20,9 @@ from litellm.llms.openai.responses.count_tokens.handler import ( ) from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, + ResponsesCountRequestParts, ) -from litellm.types.utils import LlmProviders, TokenCountResponse +from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests openai_count_tokens_handler: Final = OpenAICountTokensHandler() @@ -29,6 +37,37 @@ class OpenAITokenCounter(BaseTokenCounter): ) -> bool: return custom_llm_provider == LlmProviders.OPENAI.value + @staticmethod + def _system_instructions(system: object) -> str | None: + if system is None or isinstance(system, str): + return system + return LiteLLMAnthropicToResponsesAPIAdapter.system_blocks_to_instructions(system) or None + + @staticmethod + def _responses_request_parts( + model_to_use: str, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, + system: object, + messages_format: TokenCountMessagesFormat, + output_format: JsonValue, + ) -> ResponsesCountRequestParts: + match messages_format: + case "anthropic": + return OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model=model_to_use, messages=messages, system=system, tools=tools, output_format=output_format + ) + case "openai": + input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input( + [dict(message) for message in messages] + ) + fallback_instructions: Final = OpenAITokenCounter._system_instructions(system) + return OpenAICountTokensConfig.count_request_parts( + input_items, instructions if instructions is not None else fallback_instructions, tools + ) + case _: + assert_never(messages_format) + async def count_tokens( self, model_to_use: str, @@ -38,6 +77,8 @@ class OpenAITokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: """ Count tokens using OpenAI's Responses API /input_tokens endpoint. @@ -45,6 +86,22 @@ class OpenAITokenCounter(BaseTokenCounter): if not messages: return None + malformed_error: Final = ( + OpenAICountTokensConfig.malformed_anthropic_messages_error(messages) + if messages_format == "anthropic" + else None + ) + if malformed_error is not None: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=malformed_error, + status_code=400, + ) + deployment = deployment or {} litellm_params: Final = deployment.get("litellm_params", {}) @@ -59,12 +116,14 @@ class OpenAITokenCounter(BaseTokenCounter): api_base: Final = litellm_params.get("api_base") - # Convert chat messages to Responses API input format - input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages) - - # Use system param if instructions not extracted from messages - if instructions is None and system is not None: - instructions = system if isinstance(system, str) else str(system) + input_items, instructions, counted_tools, text = self._responses_request_parts( + model_to_use=model_to_use, + messages=messages, + tools=tools, + system=system, + messages_format=messages_format, + output_format=output_format, + ) # If no input items were produced (e.g., system-only messages), fall back to local counting if not input_items: @@ -73,11 +132,12 @@ class OpenAITokenCounter(BaseTokenCounter): try: result: Final = await openai_count_tokens_handler.handle_count_tokens_request( model=model_to_use, - input=input_items if input_items is not None else [], + input=input_items, api_key=api_key, api_base=api_base, - tools=tools, + tools=counted_tools, instructions=instructions, + text=text, ) if result is not None: diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 75684b06d05..90bfa4a8655 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -5,10 +5,18 @@ This module handles the transformation of requests to OpenAI's /v1/responses/inp """ from collections.abc import Mapping, Sequence -from typing import Any, Final, Literal +from typing import Any, Final, Literal, TypeAlias +from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict +from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.anthropic import AnthropicMessagesRequest, AnthropicOutputSchema + class ResponsesInputTextPart(TypedDict): type: ReadOnly[Literal["input_text"]] @@ -31,6 +39,30 @@ ResponsesInputPart = ResponsesInputTextPart | ResponsesInputImagePart | Response ResponsesContentRole = Literal["user", "assistant"] +ResponsesCountRequestParts: TypeAlias = tuple[ + tuple[Mapping[str, object], ...], + str | None, + tuple[Mapping[str, object], ...] | None, + Mapping[str, object] | None, +] +_RESPONSES_COUNT_REQUEST_PARTS: Final = TypeAdapter(ResponsesCountRequestParts) +_TRANSLATED_REQUEST: Final = TypeAdapter(Mapping[str, object]) + + +class _AnthropicCountMessage(TypedDict): + role: ReadOnly[str] + + +_ANTHROPIC_COUNT_MESSAGES: Final = TypeAdapter(tuple[_AnthropicCountMessage, ...]) +_ANTHROPIC_OUTPUT_SCHEMA: Final = TypeAdapter(AnthropicOutputSchema) + + +def _anthropic_output_schema(output_format: JsonValue) -> AnthropicOutputSchema | None: + try: + return _ANTHROPIC_OUTPUT_SCHEMA.validate_python(output_format) + except ValidationError: + return None + def _chat_image_block_to_responses_part(image_url: object) -> ResponsesInputImagePart | None: url: Final = image_url.get("url") if isinstance(image_url, Mapping) else image_url @@ -118,8 +150,9 @@ class OpenAICountTokensConfig: self, model: str, input: str | Sequence[object], - tools: list[dict[str, object]] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, instructions: str | None = None, + text: Mapping[str, object] | None = None, ) -> dict[str, object]: """ Transform request to OpenAI Responses API token counting format. @@ -135,7 +168,10 @@ class OpenAICountTokensConfig: request["instructions"] = instructions if tools is not None: - request["tools"] = self._transform_tools_for_responses_api(tools) + request["tools"] = self._transform_tools_for_responses_api([dict(tool) for tool in tools]) + + if text is not None: + request["text"] = text return request @@ -194,7 +230,7 @@ class OpenAICountTokensConfig: @staticmethod def messages_to_responses_input( messages: list[dict[str, Any]], - ) -> tuple: + ) -> tuple[list[dict[str, object]], str | None]: """ Convert standard chat messages format to OpenAI Responses API input format. @@ -257,3 +293,47 @@ class OpenAICountTokensConfig: instructions: Final = "\n".join(instructions_parts) if instructions_parts else None return input_items, instructions + + @staticmethod + def malformed_anthropic_messages_error(messages: Sequence[Mapping[str, object]]) -> str | None: + try: + _ANTHROPIC_COUNT_MESSAGES.validate_python(messages) + except ValidationError as e: + return "; ".join( + f"messages.{'.'.join(str(part) for part in error['loc'])}: {error['msg']}" for error in e.errors() + ) + return None + + @staticmethod + def anthropic_messages_to_responses_input( + model: str, + messages: Sequence[Mapping[str, object]], + system: object, + tools: Sequence[Mapping[str, object]] | None, + output_format: JsonValue = None, + ) -> ResponsesCountRequestParts: + """Convert with the adapter /v1/messages inference uses, so the count sees the same input it bills.""" + anthropic_request: Final = AnthropicMessagesRequest( + model=model, + messages=normalize_native_messages(messages), + system=system if isinstance(system, (str, list)) else None, + tools=[dict(tool) for tool in tools] if tools else None, + output_format=_anthropic_output_schema(output_format), + ) + translated: Final = _TRANSLATED_REQUEST.validate_python( + LiteLLMAnthropicToResponsesAPIAdapter().translate_request( + anthropic_request, include_encrypted_reasoning=False + ) + ) + return OpenAICountTokensConfig.count_request_parts( + ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(translated.get("input")), + translated.get("instructions"), + translated.get("tools"), + translated.get("text"), + ) + + @staticmethod + def count_request_parts( + input_items: object, instructions: object, tools: object, text: object = None + ) -> ResponsesCountRequestParts: + return _RESPONSES_COUNT_REQUEST_PARTS.validate_python((input_items, instructions, tools, text)) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 381986877c0..0e4cb856632 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -8,7 +8,7 @@ from types import MappingProxyType from typing import Any, Final, Literal, cast, get_type_hints import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -23,7 +23,7 @@ from litellm.types.llms.vertex_ai import ( PartType, Schema, ) -from litellm.types.utils import TokenCountResponse +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages GEMINI_VIDEO_METADATA_KEYS: Final = MappingProxyType( @@ -1278,6 +1278,8 @@ class VertexAITokenCounter(BaseTokenCounter): request_model: str = "", tools: list[dict[str, object]] | None = None, system: object | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> TokenCountResponse | None: import copy diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index d4727cc8d8d..1f17f0a42d2 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -2,10 +2,12 @@ Unified /v1/messages endpoint - (Anthropic Spec) """ +from collections.abc import Mapping from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse +from pydantic import JsonValue import litellm from litellm.anthropic_interface.exceptions import ( @@ -289,6 +291,14 @@ async def anthropic_response( ) +def _requested_output_format(data: Mapping[str, JsonValue]) -> JsonValue: + output_format: Final = data.get("output_format") + output_config: Final = data.get("output_config") + if isinstance(output_format, dict) or not isinstance(output_config, dict): + return output_format + return output_config.get("format") + + @router.post( "/v1/messages/count_tokens", tags=["[beta] Anthropic Messages Token Counting"], @@ -318,7 +328,7 @@ async def count_tokens( Returns: {"input_tokens": } """ - from litellm.proxy.proxy_server import token_counter as internal_token_counter + from litellm.proxy.proxy_server import count_request_tokens litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: @@ -345,10 +355,11 @@ async def count_tokens( system=data.get("system"), ) - # Call the internal token counter function with direct request flag set to False - token_response: Final = await internal_token_counter( + token_response: Final = await count_request_tokens( request=token_request, call_endpoint=True, + messages_format="anthropic", + output_format=_requested_output_format(data), ) _token_response_dict: dict = {} if isinstance(token_response, TokenCountResponse): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fb6a3a8a88d..621dd0d8c04 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -173,6 +173,7 @@ from litellm.types.utils import ( ModelResponseStream, StreamingChoices, TextCompletionResponse, + TokenCountMessagesFormat, TokenCountResponse, echoed_cost_map_pricing_fields, is_server_derived_pricing_key, @@ -14415,6 +14416,8 @@ async def _try_provider_token_count( request_model: str, tools: list | None = None, system: str | None = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, ) -> Optional["TokenCountResponse"]: """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting.""" if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider): @@ -14428,6 +14431,8 @@ async def _try_provider_token_count( request_model=request_model, tools=tools, system=system, + messages_format=messages_format, + output_format=output_format, ) except httpx.HTTPStatusError as e: error_message: Final = getattr(e, "message", None) or str(e) @@ -14470,8 +14475,15 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) Returns: TokenCountResponse """ - global llm_router + return await count_request_tokens(request=request, call_endpoint=call_endpoint) + +async def count_request_tokens( + request: TokenCountRequest, + call_endpoint: bool = False, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, +) -> TokenCountResponse: prompt: Final = request.prompt messages: Final = request.messages contents: Final = request.contents @@ -14530,6 +14542,8 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) request_model=request.model, tools=tools, system=system, + messages_format=messages_format, + output_format=output_format, ) if result is not None: return result diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d3c884b4a1f..a957e6957ee 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -9,6 +9,7 @@ from typing import ( Any, Final, Literal, + TypeAlias, get_args, ) @@ -4471,6 +4472,9 @@ class LiteLLMLoggingBaseClass: pass +TokenCountMessagesFormat: TypeAlias = Literal["openai", "anthropic"] + + class TokenCountResponse(LiteLLMPydanticObjectBase): total_tokens: int request_model: str diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py index bfae84e333d..3ab5826006f 100644 --- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -1,5 +1,5 @@ import json -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from typing import Final import pytest @@ -36,6 +36,146 @@ def _count_response(_request: Request) -> Reply: return Reply(body=json.dumps({"input_tokens": _COUNT}).encode()) +_PDF_BASE64: Final = "JVBERi0xLjQK" +_PNG_BASE64: Final = "aW1hZ2U=" +_LOOKUP_TOOL: Final[dict[str, JsonValue]] = { + "name": "lookup", + "description": "Find a record", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, +} +_RESPONSES_REPLY: Final[Mapping[str, JsonValue]] = { + "id": "resp_count_parity", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_count_parity", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "OK", "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}, +} +_WRAPPED_REASONING: Final = "litellm_enc:bW9kZWxfaWQ6ZGVwbG95bWVudC0x;gAAAAoriginal" +_ANSWER_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}} +_ANTHROPIC_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = { + "structured-output": { + "output_config": {"format": {"type": "json_schema", "schema": _ANSWER_SCHEMA}}, + "messages": [{"role": "user", "content": "Answer in JSON."}], + }, + "affinity-wrapped-reasoning": { + "messages": [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{_WRAPPED_REASONING}"}, + {"type": "text", "text": "OK"}, + ], + }, + {"role": "user", "content": "again"}, + ], + }, + "base64-pdf": { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the attached contract."}, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": _PDF_BASE64}, + }, + ], + } + ], + }, + "base64-image": { + "system": [{"type": "text", "text": "Describe images briefly."}], + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": _PNG_BASE64}}, + {"type": "text", "text": "What is this?"}, + ], + } + ], + }, + "replayed-web-search": { + "messages": [ + {"role": "user", "content": "What is the newest LiteLLM release?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_01", + "name": "web_search", + "input": {"query": "newest LiteLLM release"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": "LiteLLM ships a weekly release candidate and a stable line.", + } + ], + }, + {"type": "text", "text": "The newest release is in the release notes."}, + ], + }, + {"role": "user", "content": "Thanks."}, + ], + }, + "tool-turns": { + "system": "Use the lookup tool.", + "tools": [_LOOKUP_TOOL], + "messages": [ + {"role": "user", "content": "Find record 14."}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Looking it up."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "14"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}], + }, + ], + }, +} + + +def _responses_peer(request: Request) -> Reply: + if request.target == "/v1/responses/input_tokens": + return _count_response(request) + if request.target == "/v1/responses": + return Reply(body=json.dumps(_RESPONSES_REPLY).encode()) + return Reply(status=404, body=json.dumps({"error": request.target}).encode()) + + +def _body_for(wire_requests: tuple[Request, ...], target: str) -> dict[str, JsonValue]: + matching: Final = tuple(request for request in wire_requests if request.target == target) + assert len(matching) == 1, [request.target for request in wire_requests] + return _JSON_OBJECT.validate_json(matching[0].body) + + +def _counted_fields(body: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + return {key: body[key] for key in ("input", "instructions", "tools", "text") if body.get(key) is not None} + + def _count_body(wire: Wire) -> dict[str, JsonValue]: requests: Final = wire.drain() assert len(requests) == 1, requests @@ -159,3 +299,53 @@ def test_utils_token_counter_keeps_chat_tool_conversion(counting_proxy: OwnedPro "parameters": _TOOL_SCHEMA, } ] + + +@pytest.mark.parametrize("case", sorted(_ANTHROPIC_BODIES)) +def test_messages_count_tokens_sends_the_input_v1_messages_inference_sends( + counting_proxy: OwnedProxy, case: str +) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_responses_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + body: Final = {"model": model, **_ANTHROPIC_BODIES[case]} + generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16}) + counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body) + assert generated.status_code == 200, generated.text + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT + requests: Final = wire.drain() + inference_body: Final = _body_for(requests, "/v1/responses") + count_body: Final = _body_for(requests, "/v1/responses/input_tokens") + assert _counted_fields(count_body) == _counted_fields(inference_body) + assert count_body["model"] == inference_body["model"] + + +def test_messages_count_tokens_without_a_role_answers_400_before_calling_openai(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": [{"role": "user", "content": "hello"}, {"content": "no role"}]}, + ) + assert response.status_code == 400, response.text + assert "messages.1.role: Field required" in response.text + assert wire.drain() == () + + +def test_google_count_tokens_sends_the_system_instruction_text(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_count_response) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:countTokens", + { + "systemInstruction": {"parts": [{"text": "Be brief."}, {"text": "Answer in English."}]}, + "contents": [{"role": "user", "parts": [{"text": "hello"}]}], + }, + ) + assert response.status_code == 200, response.text + assert _count_body(wire)["instructions"] == "Be brief.\nAnswer in English." diff --git a/tests/integration/spend/test_cache_and_quota.py b/tests/integration/spend/test_cache_and_quota.py index d02ebccf710..801a14c87d3 100644 --- a/tests/integration/spend/test_cache_and_quota.py +++ b/tests/integration/spend/test_cache_and_quota.py @@ -369,7 +369,12 @@ def test_repeated_count_tokens_on_budgeted_key_does_not_reserve_budget_or_block_ { "path": f"/{counting_upstream.scenario_id}/v1/responses/input_tokens", "authorization": "Bearer integration-provider-key", - "body": {"model": "gpt-4o-mini", "input": [{"role": "user", "content": "hello!!!"}]}, + "body": { + "model": "gpt-4o-mini", + "input": [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hello!!!"}]} + ], + }, "method": "POST", "api_key": "", } diff --git a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py index 2e0fc55cda4..976ecff01ea 100644 --- a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py +++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py @@ -5,11 +5,12 @@ from typing import Final import httpx import pytest import respx -from pydantic import ValidationError +from pydantic import JsonValue, ValidationError import litellm from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter -from litellm.types.utils import TokenCountResponse +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse COUNT_URL: Final = "https://api.openai.com/v1/responses/input_tokens" @@ -163,6 +164,208 @@ async def test_upstream_error_status_becomes_an_error_response_with_that_status_ ) +async def _sent_count_body( + monkeypatch: pytest.MonkeyPatch, + messages: list[dict[str, object]], + system: object, + messages_format: TokenCountMessagesFormat, + output_format: JsonValue = None, +) -> object: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1})) + await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=messages, + contents=None, + deployment={"litellm_params": {"api_key": "sk-test"}}, + request_model="openai/gpt-4o", + system=system, + messages_format=messages_format, + output_format=output_format, + ) + return json.loads(route.calls.last.request.content) + + +@pytest.mark.asyncio +async def test_anthropic_format_sends_the_document_image_and_tool_turns_to_the_count_endpoint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + messages: Final[list[dict[str, object]]] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize."}, + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBE"}}, + {"type": "image", "source": {"type": "url", "url": "https://example.com/page.png"}}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "found"}]}, + ] + + sent: Final = await _sent_count_body( + monkeypatch, messages, [{"type": "text", "text": "Be brief."}], messages_format="anthropic" + ) + + assert sent == { + "model": "gpt-4o", + "input": [ + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize."}, + {"type": "input_file", "filename": "document.pdf", "file_data": "data:application/pdf;base64,JVBE"}, + {"type": "input_image", "image_url": "https://example.com/page.png", "detail": "auto"}, + ], + }, + {"type": "function_call", "call_id": "toolu_01", "name": "lookup", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "toolu_01", "output": "found"}, + ], + "instructions": "Be brief.", + } + + +@pytest.mark.asyncio +async def test_openai_format_keeps_the_chat_conversion(monkeypatch: pytest.MonkeyPatch) -> None: + messages: Final[list[dict[str, object]]] = [ + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": [{"type": "text", "text": "Summarize."}]}, + {"role": "tool", "tool_call_id": "call_1", "content": "found"}, + ] + + sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="openai") + + assert sent == { + "model": "gpt-4o", + "input": [ + {"role": "user", "content": "Summarize."}, + {"type": "function_call_output", "call_id": "call_1", "output": "found"}, + ], + "instructions": "Be brief.", + } + + +@pytest.mark.asyncio +async def test_anthropic_format_sends_affinity_wrapped_reasoning_unwrapped_as_inference_does( + monkeypatch: pytest.MonkeyPatch, +) -> None: + wrapped: Final = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("gAAAAoriginal", "deployment-1") + messages: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"}, + {"type": "text", "text": "OK"}, + ], + }, + {"role": "user", "content": "again"}, + ] + + sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic") + + assert sent == { + "model": "gpt-4o", + "input": [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + {"type": "reasoning", "summary": [], "encrypted_content": "gAAAAoriginal"}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "OK"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "again"}]}, + ], + } + + +@pytest.mark.asyncio +async def test_anthropic_format_counts_the_structured_output_schema(monkeypatch: pytest.MonkeyPatch) -> None: + schema: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}} + + sent: Final = await _sent_count_body( + monkeypatch, + [{"role": "user", "content": "Answer in JSON."}], + system=None, + messages_format="anthropic", + output_format={"type": "json_schema", "schema": schema}, + ) + + assert sent == { + "model": "gpt-4o", + "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Answer in JSON."}]}], + "text": {"format": {"type": "json_schema", "name": "structured_output", "schema": schema, "strict": False}}, + } + + +@pytest.mark.asyncio +async def test_anthropic_format_sends_a_replayed_web_search_result_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + messages: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + {"role": "user", "content": "Thanks."}, + ] + + sent: Final = json.dumps(await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic")) + + assert snippet in sent + assert "https://docs.litellm.ai/release_notes" in sent + + +@pytest.mark.asyncio +async def test_anthropic_format_message_without_a_role_answers_400_without_calling_openai( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1})) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}, {"content": "no role"}], + contents=None, + deployment={"litellm_params": {"api_key": "sk-test"}}, + request_model="openai/gpt-4o", + messages_format="anthropic", + ) + + assert route.called is False + assert result is not None + assert (result.error, result.status_code, result.error_message) == (True, 400, "messages.1.role: Field required") + + +@pytest.mark.asyncio +async def test_openai_format_sends_a_list_system_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None: + sent: Final = await _sent_count_body( + monkeypatch, + [{"role": "user", "content": "hi"}], + system=[{"type": "text", "text": "Be brief."}, {"type": "text", "text": "Answer in English."}], + messages_format="openai", + ) + + assert sent == { + "model": "gpt-4o", + "input": [{"role": "user", "content": "hi"}], + "instructions": "Be brief.\nAnswer in English.", + } + + @pytest.mark.asyncio async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py index 98ddfa08018..0a5039791ee 100644 --- a/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_count_tokens_transformation.py @@ -2,6 +2,9 @@ from typing import Final import pytest +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, ) @@ -484,3 +487,108 @@ def test_get_required_headers(): assert headers["Authorization"] == "Bearer sk-test-key" assert headers["Content-Type"] == "application/json" + + +PDF_BASE64: Final = "JVBERi0xLjcK" +PNG_BASE64: Final = "iVBORw0KGgo=" + +ANTHROPIC_AGENT_TURN: Final[list[dict[str, object]]] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the indemnity terms."}, + { + "type": "document", + "title": "contract.pdf", + "source": {"type": "base64", "media_type": "application/pdf", "data": PDF_BASE64}, + }, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PNG_BASE64}}, + ], + }, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me pull clause 14."}, + {"type": "tool_use", "id": "toolu_01", "name": "lookup_clause", "input": {"id": "14"}}, + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}], + }, +] + +ANTHROPIC_TOOLS: Final[list[dict[str, object]]] = [ + { + "name": "lookup_clause", + "description": "Fetch a clause by id", + "input_schema": {"type": "object", "properties": {"id": {"type": "string"}}}, + } +] + + +@pytest.mark.parametrize( + "system", + [None, "You review contracts.", [{"type": "text", "text": "You review contracts."}]], +) +def test_anthropic_messages_convert_exactly_as_v1_messages_inference_converts_them(system: object) -> None: + inference_request: Final[dict[str, object]] = { + "model": "gpt-4o", + "messages": ANTHROPIC_AGENT_TURN, + "tools": ANTHROPIC_TOOLS, + **({"system": system} if system is not None else {}), + } + inference: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_request( + inference_request, include_encrypted_reasoning=False + ) + + counted: Final = OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=system, tools=ANTHROPIC_TOOLS + ) + + assert counted == ( + tuple(inference["input"]), + inference.get("instructions"), + tuple(inference["tools"]), + inference.get("text"), + ) + + +def test_anthropic_messages_keep_every_block_the_chat_converter_drops() -> None: + input_items, _, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=None, tools=None + ) + + assert input_items == ( + { + "type": "message", + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize the indemnity terms."}, + { + "type": "input_file", + "filename": "contract.pdf", + "file_data": f"data:application/pdf;base64,{PDF_BASE64}", + }, + {"type": "input_image", "image_url": f"data:image/png;base64,{PNG_BASE64}", "detail": "auto"}, + ], + }, + {"type": "function_call", "call_id": "toolu_01", "name": "lookup_clause", "arguments": '{"id": "14"}'}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Let me pull clause 14."}], + }, + {"type": "function_call_output", "call_id": "toolu_01", "output": "14. Indemnification."}, + ) + + +def test_anthropic_list_system_becomes_joined_instruction_text() -> None: + _, instructions, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + system=[{"type": "text", "text": "First rule."}, {"type": "text", "text": "Second rule."}], + tools=None, + ) + + assert instructions == "First rule.\nSecond rule." diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index c0e7a4fdb47..69d182667ec 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -278,13 +278,17 @@ class TestHttpExceptionDictDetail: request.headers = {} with ( - patch.object(ep, "read_request_body", new=AsyncMock(return_value={})), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object( + ep, "read_request_body", new=AsyncMock(return_value={}) + ), # test-quality-ok: endpoint reads the body via a module function; no injection seam patch.object( # test-quality-ok: the guardrail raise happens deep inside this call; the test targets the endpoint's except block ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=AsyncMock(side_effect=exc), ), - patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam + patch.object( + proxy_server, "proxy_logging_obj" + ) as mock_logging, # test-quality-ok: module global imported at call time; no injection seam ): mock_logging.post_call_failure_hook = AsyncMock() response = await ep.anthropic_response( @@ -375,9 +379,15 @@ class TestErrorLogCarriesCallId: request.headers = {} with ( - patch.object(ep, "read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})), # test-quality-ok: endpoint reads the body via a module function; no injection seam - patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process), # test-quality-ok: the provider failure happens inside this call; the test targets the endpoint's except block - patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam + patch.object( + ep, "read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"}) + ), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object( + ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process + ), # test-quality-ok: the provider failure happens inside this call; the test targets the endpoint's except block + patch.object( + proxy_server, "proxy_logging_obj" + ) as mock_logging, # test-quality-ok: module global imported at call time; no injection seam caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), ): mock_logging.post_call_failure_hook = AsyncMock() @@ -408,9 +418,15 @@ class TestErrorLogCarriesCallId: request.headers = {} with ( - patch.object(ep, "read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})), # test-quality-ok: endpoint reads the body via a module function; no injection seam - patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process), # test-quality-ok: the proxy shaped failure happens inside this call; the test targets the endpoint's except block - patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam + patch.object( + ep, "read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"}) + ), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object( + ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process + ), # test-quality-ok: the proxy shaped failure happens inside this call; the test targets the endpoint's except block + patch.object( + proxy_server, "proxy_logging_obj" + ) as mock_logging, # test-quality-ok: module global imported at call time; no injection seam ): mock_logging.post_call_failure_hook = AsyncMock() response = await ep.anthropic_response( @@ -440,7 +456,9 @@ class TestErrorLogCarriesCallId: "read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet", "messages": [{"role": "user", "content": "hi"}]}), ), - patch.object(proxy_server, "token_counter", new=AsyncMock(side_effect=RuntimeError("tokenizer down"))), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block + patch.object( + proxy_server, "count_request_tokens", new=AsyncMock(side_effect=RuntimeError("tokenizer down")) + ), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(HTTPException) as raised, ): diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index c93c199cfa2..8cd6887d891 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -30,7 +30,6 @@ from litellm.proxy.proxy_server import token_counter from litellm.types.utils import TokenCountResponse - @pytest.mark.asyncio async def test_vLLM_token_counting(): """ @@ -63,9 +62,7 @@ async def test_vLLM_token_counting(): print("response: ", response) - assert ( - response.tokenizer_type == "openai_tokenizer" - ) # SHOULD use the default tokenizer + assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the default tokenizer assert response.model_used == "wolfram/miquliz-120b-v2.0" @@ -98,9 +95,7 @@ async def test_token_counting_model_not_in_model_list(): print("response: ", response) - assert ( - response.tokenizer_type == "openai_tokenizer" - ) # SHOULD use the OpenAI tokenizer + assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the OpenAI tokenizer assert response.model_used == "special-alias" @@ -133,9 +128,7 @@ async def test_gpt_token_counting(): print("response: ", response) - assert ( - response.tokenizer_type == "openai_tokenizer" - ) # SHOULD use the OpenAI tokenizer + assert response.tokenizer_type == "openai_tokenizer" # SHOULD use the OpenAI tokenizer assert response.request_model == "gpt-4" @@ -173,10 +166,8 @@ async def test_anthropic_messages_count_tokens_endpoint(): anthropic_endpoints.read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=False): - assert ( - call_endpoint == True - ), "Should be called with call_endpoint=True for Anthropic endpoint" + async def mock_token_counter(request, call_endpoint=False, messages_format="openai", output_format=None): + assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "claude-3-sonnet-20240229" assert request.messages == [{"role": "user", "content": "Hello Claude!"}] @@ -192,8 +183,8 @@ async def test_anthropic_messages_count_tokens_endpoint(): # Patch the imported token_counter function from proxy_server import litellm.proxy.proxy_server as proxy_server - original_token_counter = proxy_server.token_counter - proxy_server.token_counter = mock_token_counter + original_token_counter = proxy_server.count_request_tokens + proxy_server.count_request_tokens = mock_token_counter try: # Call the endpoint @@ -210,7 +201,7 @@ async def test_anthropic_messages_count_tokens_endpoint(): finally: # Restore original functions anthropic_endpoints.read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -247,10 +238,8 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): anthropic_endpoints.read_request_body = mock_read_request_body # Mock the internal token_counter function to return a controlled response - async def mock_token_counter(request, call_endpoint=True): - assert ( - call_endpoint == True - ), "Should be called with call_endpoint=True for Anthropic endpoint" + async def mock_token_counter(request, call_endpoint=True, messages_format="openai", output_format=None): + assert call_endpoint == True, "Should be called with call_endpoint=True for Anthropic endpoint" assert request.model == "gpt-4" assert request.messages == [{"role": "user", "content": "Hello GPT!"}] @@ -266,8 +255,8 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): # Patch the imported token_counter function from proxy_server import litellm.proxy.proxy_server as proxy_server - original_token_counter = proxy_server.token_counter - proxy_server.token_counter = mock_token_counter + original_token_counter = proxy_server.count_request_tokens + proxy_server.count_request_tokens = mock_token_counter try: # Call the endpoint @@ -284,7 +273,7 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): finally: # Restore original functions anthropic_endpoints.read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -407,9 +396,7 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter(): # Mock the global handler instance in token_counter module mock_handler = MagicMock() - mock_handler.handle_count_tokens_request = AsyncMock( - return_value={"input_tokens": 42} - ) + mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42}) with patch( "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", @@ -466,9 +453,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): # Mock the global handler instance in token_counter module mock_handler = MagicMock() - mock_handler.handle_count_tokens_request = AsyncMock( - return_value={"input_tokens": 42} - ) + mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42}) with patch( "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", @@ -527,9 +512,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): # Mock the global handler instance in token_counter module mock_handler = MagicMock() - mock_handler.handle_count_tokens_request = AsyncMock( - return_value={"input_tokens": 42} - ) + mock_handler.handle_count_tokens_request = AsyncMock(return_value={"input_tokens": 42}) with patch( "litellm.llms.anthropic.count_tokens.token_counter.anthropic_count_tokens_handler", @@ -544,9 +527,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): mock_router.model_list = [ { "model_name": "claude-3-5-sonnet", - "litellm_params": { - "model": "anthropic/claude-3-5-sonnet-20241022" - }, + "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}, "model_info": {}, } ] @@ -555,9 +536,7 @@ async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): mock_router.async_get_available_deployment = AsyncMock( return_value={ "model_name": "claude-3-5-sonnet", - "litellm_params": { - "model": "anthropic/claude-3-5-sonnet-20241022" - }, + "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}, "model_info": {}, } ) @@ -592,9 +571,7 @@ async def test_factory_registration(): assert counter is not None # Create test deployments - anthropic_deployment = { - "litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"} - } + anthropic_deployment = {"litellm_params": {"model": "anthropic/claude-3-5-sonnet-20241022"}} non_anthropic_deployment = {"litellm_params": {"model": "openai/gpt-4"}} @@ -630,9 +607,7 @@ async def test_bedrock_count_tokens_endpoint(): model_list=[ { "model_name": "claude-bedrock", - "litellm_params": { - "model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0" - }, + "litellm_params": {"model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0"}, } ] ) @@ -646,9 +621,7 @@ async def test_bedrock_count_tokens_endpoint(): } # Test the mock handler directly to verify correct parameter extraction - await mock_count_tokens_handler( - request_data, {}, "anthropic.claude-3-sonnet-20240229-v1:0" - ) + await mock_count_tokens_handler(request_data, {}, "anthropic.claude-3-sonnet-20240229-v1:0") @pytest.mark.asyncio @@ -710,10 +683,7 @@ async def test_vertex_ai_anthropic_token_counting(): assert call_args is not None assert call_args.kwargs["model"] == "claude-3-5-sonnet-20241022" assert "messages" in call_args.kwargs["request_data"] - assert ( - call_args.kwargs["request_data"]["messages"][0]["content"] - == "Hello Claude on Vertex AI! How are you?" - ) + assert call_args.kwargs["request_data"]["messages"][0]["content"] == "Hello Claude on Vertex AI! How are you?" # Validate response structure assert response.model_used == "claude-3-5-sonnet-20241022" @@ -746,9 +716,7 @@ def test_vertex_ai_partner_models_token_counting_endpoint(vertex_location): if vertex_location == "global": assert endpoint.startswith("https://aiplatform.googleapis.com") else: - assert endpoint.startswith( - f"https://{vertex_location}-aiplatform.googleapis.com" - ) + assert endpoint.startswith(f"https://{vertex_location}-aiplatform.googleapis.com") class _FailingRuntimeHandler(BedrockCountTokensHandler): @@ -834,20 +802,14 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): with patch.object(handler, "validate_count_tokens_request"): with patch.object(handler, "_get_aws_region_name", return_value="us-west-2"): - with patch.object( - handler, "transform_anthropic_to_bedrock_count_tokens", return_value={} - ): + with patch.object(handler, "transform_anthropic_to_bedrock_count_tokens", return_value={}): with patch.object( handler, "get_bedrock_count_tokens_endpoint", return_value="https://example.com", ): - with patch.object( - handler, "_sign_request", return_value=({}, "{}") - ): - with patch( - "litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client" - ) as mock_client: + with patch.object(handler, "_sign_request", return_value=({}, "{}")): + with patch("litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client") as mock_client: mock_async_client = AsyncMock() mock_async_client.post = AsyncMock(side_effect=http_error) mock_client.return_value = mock_async_client @@ -856,9 +818,7 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): await handler.handle_count_tokens_request( request_data={ "model": "test", - "messages": [ - {"role": "user", "content": "hello"} - ], + "messages": [{"role": "user", "content": "hello"}], }, litellm_params={}, resolved_model="anthropic.claude-3-sonnet", @@ -866,10 +826,7 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): assert exc_info.value.status_code == 403 # Message should be the raw response text - assert ( - exc_info.value.message - == "Forbidden - Invalid credentials" - ) + assert exc_info.value.message == "Forbidden - Invalid credentials" @pytest.mark.asyncio @@ -897,9 +854,7 @@ async def test_token_counter_httpx_status_error_raises_proxy_exception(): mock_counter.count_tokens = AsyncMock(side_effect=http_error) # Save originals - original_get_provider_token_counter = ( - litellm.proxy.proxy_server._get_provider_token_counter - ) + original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter original_router = litellm.proxy.proxy_server.llm_router try: @@ -907,9 +862,7 @@ async def test_token_counter_httpx_status_error_raises_proxy_exception(): def mock_get_provider_token_counter(deployment, model_to_use): return (mock_counter, "claude-4-6-sonnet", "vertex_ai") - litellm.proxy.proxy_server._get_provider_token_counter = ( - mock_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter mock_router = MagicMock() mock_router.async_get_available_deployment = AsyncMock( @@ -937,9 +890,7 @@ async def test_token_counter_httpx_status_error_raises_proxy_exception(): assert exc_info.value.type == "token_counting_error" assert exc_info.value.param == "model" finally: - litellm.proxy.proxy_server._get_provider_token_counter = ( - original_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter litellm.proxy.proxy_server.llm_router = original_router @@ -975,9 +926,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): # Save original value and function original_disable = litellm.disable_token_counter - original_get_provider_token_counter = ( - litellm.proxy.proxy_server._get_provider_token_counter - ) + original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter try: litellm.disable_token_counter = True @@ -991,9 +940,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): def mock_get_provider_token_counter(deployment, model_to_use): return (mock_counter, "anthropic.claude-3-sonnet", "bedrock") - litellm.proxy.proxy_server._get_provider_token_counter = ( - mock_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter with pytest.raises(ProxyException) as exc_info: await token_counter( @@ -1008,9 +955,7 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled(): assert "Rate limit exceeded" in exc_info.value.message finally: litellm.disable_token_counter = original_disable - litellm.proxy.proxy_server._get_provider_token_counter = ( - original_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter @pytest.mark.asyncio @@ -1045,9 +990,7 @@ async def test_proxy_token_counter_error_raises_when_enabled(): # Save original value and function original_disable = litellm.disable_token_counter - original_get_provider_token_counter = ( - litellm.proxy.proxy_server._get_provider_token_counter - ) + original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter try: litellm.disable_token_counter = False @@ -1061,9 +1004,7 @@ async def test_proxy_token_counter_error_raises_when_enabled(): def mock_get_provider_token_counter(deployment, model_to_use): return (mock_counter, "anthropic.claude-3-sonnet", "bedrock") - litellm.proxy.proxy_server._get_provider_token_counter = ( - mock_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter with patch("litellm.token_counter") as mock_local_counter: with pytest.raises(ProxyException) as exc_info: @@ -1081,9 +1022,7 @@ async def test_proxy_token_counter_error_raises_when_enabled(): mock_local_counter.assert_not_called() finally: litellm.disable_token_counter = original_disable - litellm.proxy.proxy_server._get_provider_token_counter = ( - original_get_provider_token_counter - ) + litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter @pytest.mark.asyncio @@ -1109,10 +1048,10 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): original_read_request_body = anthropic_endpoints.read_request_body anthropic_endpoints.read_request_body = mock_read_request_body - original_token_counter = proxy_server.token_counter + original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with Bedrock-style error - async def mock_token_counter_error(request, call_endpoint=False): + async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): raise ProxyException( message='{"detail":{"message":"Input is too long for requested model."}}', type="token_counting_error", @@ -1120,7 +1059,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): code=400, ) - proxy_server.token_counter = mock_token_counter_error + proxy_server.count_request_tokens = mock_token_counter_error try: with pytest.raises(HTTPException) as exc_info: @@ -1136,7 +1075,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): assert detail["error"]["message"] == "Input is too long for requested model." finally: anthropic_endpoints.read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -1161,10 +1100,10 @@ async def test_anthropic_endpoint_403_permission_error_format(): original_read_request_body = anthropic_endpoints.read_request_body anthropic_endpoints.read_request_body = mock_read_request_body - original_token_counter = proxy_server.token_counter + original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 403 error - async def mock_token_counter_error(request, call_endpoint=False): + async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): raise ProxyException( message='{"Message":"Bearer Token has expired"}', type="token_counting_error", @@ -1172,7 +1111,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): code=403, ) - proxy_server.token_counter = mock_token_counter_error + proxy_server.count_request_tokens = mock_token_counter_error try: with pytest.raises(HTTPException) as exc_info: @@ -1186,7 +1125,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): assert detail["error"]["message"] == "Bearer Token has expired" finally: anthropic_endpoints.read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter @pytest.mark.asyncio @@ -1211,10 +1150,10 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): original_read_request_body = anthropic_endpoints.read_request_body anthropic_endpoints.read_request_body = mock_read_request_body - original_token_counter = proxy_server.token_counter + original_token_counter = proxy_server.count_request_tokens # Mock token_counter to raise ProxyException with 429 error - async def mock_token_counter_error(request, call_endpoint=False): + async def mock_token_counter_error(request, call_endpoint=False, messages_format="openai", output_format=None): raise ProxyException( message="Rate limit exceeded", type="token_counting_error", @@ -1222,7 +1161,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): code=429, ) - proxy_server.token_counter = mock_token_counter_error + proxy_server.count_request_tokens = mock_token_counter_error try: with pytest.raises(HTTPException) as exc_info: @@ -1236,4 +1175,4 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): assert detail["error"]["message"] == "Rate limit exceeded" finally: anthropic_endpoints.read_request_body = original_read_request_body - proxy_server.token_counter = original_token_counter + proxy_server.count_request_tokens = original_token_counter diff --git a/tests/unit/test_count_tokens_public_api.py b/tests/unit/test_count_tokens_public_api.py index 84121dd63e4..e33bc1e5907 100644 --- a/tests/unit/test_count_tokens_public_api.py +++ b/tests/unit/test_count_tokens_public_api.py @@ -97,7 +97,7 @@ def test_acount_tokens_with_tools(): assert result.total_tokens == 30 mock_handler.assert_called_once() call_kwargs = mock_handler.call_args - assert call_kwargs.kwargs.get("tools") == tools + assert list(call_kwargs.kwargs.get("tools") or ()) == tools def test_acount_tokens_with_system():