diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index 7a7d1d776d3..46c90295198 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -125,6 +125,8 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): status_code=e.response.status_code, message=e.response.text, ) + except litellm.Timeout as e: + raise AnthropicError(status_code=e.status_code, message=str(e)) from e except Exception as e: verbose_logger.error("Error in CountTokens handler: %s", e) raise AnthropicError( diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 99b1cd7f241..755ccd129d8 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -18,7 +18,7 @@ from litellm.llms.anthropic.count_tokens.transformation import ( ChatCountTranslationError, anthropic_count_fields, ) -from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.llms.base_llm.base_utils import BaseTokenCounter, deployment_timeout from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests @@ -134,6 +134,7 @@ class AnthropicTokenCounter(BaseTokenCounter): optional_params=count_options, context_management=context_management.spec if context_management is not None else None, drop_params=drops_unsupported_params(litellm_params.get("drop_params")), + timeout=deployment_timeout(litellm_params), ) if result is not None: diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 5fcb84593d2..443a3a5edad 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -147,6 +147,8 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): status_code=e.response.status_code, message=e.response.text, ) + except litellm.Timeout as e: + raise AnthropicError(status_code=e.status_code, message=str(e)) from e except Exception as e: verbose_logger.error("Error in CountTokens handler: %s", e) raise AnthropicError( 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 8657fcddab4..5031b69f83b 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -20,7 +20,7 @@ from litellm.llms.anthropic.count_tokens.transformation import ( from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, ) -from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.llms.base_llm.base_utils import BaseTokenCounter, deployment_timeout from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests @@ -137,6 +137,7 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): system=system, optional_params=count_options, context_management=None if context_management is None else context_management.spec, + timeout=deployment_timeout(litellm_params), ) if result is not None: diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index 45f623730a8..8b1ccda3a61 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -15,8 +15,9 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Any, Final, TypeAlias +import httpx from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # SDK parser internals -from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, JsonValue, PositiveFloat, TypeAdapter, ValidationError from typing_extensions import TypeIs # noqa: TID251 # narrows a message to its system member by the role literal from litellm._logging import verbose_logger @@ -38,6 +39,18 @@ from litellm.types.utils import ( TokenCountResponse, ) +_SECONDS: Final[TypeAdapter[float]] = TypeAdapter(PositiveFloat) + + +def deployment_timeout(litellm_params: Mapping[str, object]) -> float | httpx.Timeout | None: + timeout: Final = litellm_params.get("timeout") + if timeout is None or isinstance(timeout, httpx.Timeout): + return timeout + try: + return _SECONDS.validate_python(timeout) + except ValidationError: + return None + class BaseTokenCounter(ABC): @abstractmethod diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 2cf48a75524..ce9ed6c0940 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -111,6 +111,8 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): status_code=e.response.status_code, message=e.response.text, ) + except litellm.Timeout as e: + raise OpenAIError(status_code=e.status_code, message=str(e)) from e except (httpx.RequestError, json.JSONDecodeError, ValueError) as e: verbose_logger.error("Error in CountTokens handler: %s", e) raise OpenAIError( diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index c82a2982a54..cd60bb1327f 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -7,6 +7,7 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Any, Final +import httpx import openai from pydantic import JsonValue from typing_extensions import assert_never @@ -21,7 +22,7 @@ from litellm.llms.anthropic.messages.responses_adapters.handler import count_vie from litellm.llms.anthropic.messages.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, ) -from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.llms.base_llm.base_utils import BaseTokenCounter, deployment_timeout from litellm.llms.openai.common_utils import OpenAIError, decoded_json_object, get_openai_credentials from litellm.llms.openai.responses.count_tokens.handler import ( OpenAICountTokensHandler, @@ -61,6 +62,7 @@ def _workload_identity_config( class _CountCredential: api_base: str | None credential: str | OpenAIWorkloadIdentityConfig + timeout: float | httpx.Timeout | None def _count_credential(deployment: Mapping[str, object] | None, model_label: str) -> _CountCredential | OpenAIError: @@ -82,7 +84,9 @@ def _count_credential(deployment: Mapping[str, object] | None, model_label: str) "workload identity on the deployment, or OPENAI_API_KEY in the proxy environment." ), ) - return _CountCredential(api_base=credentials.api_base, credential=credential) + return _CountCredential( + api_base=credentials.api_base, credential=credential, timeout=deployment_timeout(litellm_params) + ) def _count_error(model_to_use: str, request_model: str, message: str, status_code: int) -> TokenCountResponse: @@ -245,6 +249,7 @@ class OpenAITokenCounter(BaseTokenCounter): tools=counted_tools, instructions=instructions, text=text, + timeout=resolved.timeout, ), model_to_use, request_model, @@ -264,7 +269,9 @@ class OpenAITokenCounter(BaseTokenCounter): body: Final = native_count_request(request, model_to_use, policy) return await self._counted( resolved, - lambda api_key, api_base: openai_count_tokens_handler.post_count_request(body, api_key, api_base), + lambda api_key, api_base: openai_count_tokens_handler.post_count_request( + body, api_key, api_base, resolved.timeout + ), model_to_use, request_model, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e6447754436..3fddca6f4d8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15325,7 +15325,12 @@ def _provider_count_outcome( if result is not None and result.error is True: return ProviderCountFailure( status_code=result.status_code or 500, - message=result.error_message or "Token counting failed", + message=result.error_message + or ( + "Token counting failed" + if result.status_code is None + else f"The provider's count API answered {result.status_code} with no error message" + ), rejected_before_provider_call=result.rejected_before_provider_call, ) return result @@ -15794,6 +15799,14 @@ def _token_count_settings_model(deployment: _TokenCountDeployment, request_model _PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset( {litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value} ) +_CHAT_MODES: Final = frozenset({"chat", "responses"}) + + +def _serves_chat(deployment: _TokenCountDeployment, cost_map_row: Mapping[str, object]) -> bool: + mode: Final = _non_empty_str(deployment.get("model_info", {}).get("mode")) or _non_empty_str( + cost_map_row.get("mode") + ) + return mode is None or mode in _CHAT_MODES _ESTIMABLE_GOOGLE_PART_KEYS: Final = frozenset({"text", "functionCall", "functionResponse", "thoughtSignature"}) @@ -16104,6 +16117,12 @@ async def _routed_count_deployment( return None +def _with_chat_timeout(litellm_params: Mapping[str, object]) -> Mapping[str, object]: + if llm_router is None: + return litellm_params + return {**litellm_params, "timeout": llm_router.settings().timeout_for(litellm_params)} + + async def _count_request_tokens( request: TokenCountRequest, routed: Mapping[str, object] | None, @@ -16128,8 +16147,8 @@ async def _count_request_tokens( litellm_model_name = deployment.get("litellm_params", {}).get("model") deployment = { # rebind-ok: merge the named credential into a copy **deployment, - "litellm_params": with_named_credential( - _LITELLM_PARAMS_COPY.validate_python(deployment.get("litellm_params", {})) + "litellm_params": _with_chat_timeout( + with_named_credential(_LITELLM_PARAMS_COPY.validate_python(deployment.get("litellm_params", {}))) ), } # remove the custom_llm_provider_prefix in the litellm_model_name @@ -16142,9 +16161,10 @@ async def _count_request_tokens( deployment_view: Final = _TOKEN_COUNT_DEPLOYMENT.validate_python( deployment or _model_group_deployment(llm_router, request.model) or {} ) + cost_map_row: Final = await asyncify(cost_map_entry)(*_token_count_settings_model(deployment_view, request.model)) settings: Final = await asyncify(resolve_token_count_settings)( configured=deployment_view.get("model_info", {}), - default=await asyncify(cost_map_entry)(*_token_count_settings_model(deployment_view, request.model)), + default=cost_map_row, deployment_region=partial(_deployment_aws_region, deployment_view), ) @@ -16169,6 +16189,15 @@ async def _count_request_tokens( code=400, ) prompt_only: Final = bool(prompt) and conversation_empty + prompt_as_turn: Final = ( + prompt_only + and provider_counter is not None + and provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider) + and _serves_chat(deployment_view, cost_map_row) + ) + prompt_as_message: Final = prompt_as_turn or ( + prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES + ) request_messages: Final = _counted_request_messages(messages, messages_format, count_deployment) local_count_error: Final = malformed_tool_result_count_error(request_messages, messages_format, count_deployment) malformed_messages: Final = ( @@ -16205,7 +16234,6 @@ async def _count_request_tokens( ) inference_system: Final = system if edited_request is None else edited_request.system contents_estimable: Final = not contents or (bool(messages) and _estimable_contents(contents)) - prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES routed_deployment: Final = _ROUTED_DEPLOYMENT.validate_python(deployment) routed_request: Final = ( None if responses_request is None else with_provider_file_ids(responses_request, routed_deployment) @@ -16306,7 +16334,7 @@ async def _count_request_tokens( _with_get(model_info).get("custom_tokenizer", None), ) system_message: Final = _system_message(local_system) - prompt_turn: Final = prompt_only and not (system_message is None and not tools) + prompt_turn: Final = prompt_as_turn or (prompt_only and not (system_message is None and not tools)) typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes Sequence[AllMessageValues] | None, [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else local_messages), diff --git a/litellm/router_api.py b/litellm/router_api.py index d07b135e9a5..c85a764151f 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -1308,6 +1308,9 @@ class RouterSettings: """The guardrail configured under this name (was ``get_available_guardrail``).""" return self._owner("_core_get_available_guardrail")._core_get_available_guardrail(name) + def timeout_for(self, litellm_params: Mapping[str, object]) -> object: + return self._owner("_core_deployment_call_timeout")._core_deployment_call_timeout(litellm_params) + class RouterHealth: """Cooldowns, blocked deployments, health checks and usage (``health()``).""" diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index b7e67fffd57..682aa557007 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -113,7 +113,7 @@ from litellm.router_rust.invoke.endpoints import ( stamp_error, unrouted_fields, ) -from litellm.router_rust.invoke.shaping import Shaping +from litellm.router_rust.invoke.shaping import Shaping, non_stream_timeout from litellm.router_rust.invoke.slots import enter_deployment_slot, max_parallel_requests_limit from litellm.router_rust.invoke.streaming import ( StreamAttempt, @@ -1592,6 +1592,9 @@ class RustCore( def _shaping(self) -> Shaping: return self._endpoints.shaping(self._resolved) + def _core_deployment_call_timeout(self, litellm_params: Mapping[str, object]) -> object: + return non_stream_timeout(self._shaping, MappingProxyType({}), litellm_params) + # -- building -------------------------------------------------------------------------------- def _new_handle(self, model_list: Sequence[object]) -> "RouterCoreHandle": diff --git a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py index 380881a6796..86a8cb3df4f 100644 --- a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py +++ b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py @@ -332,6 +332,22 @@ def test_utils_token_counter_local_mode_never_calls_the_peer_for_a_leading_syste assert peer.drain() == () +def test_utils_token_counter_call_endpoint_counts_a_prompt_as_one_user_message( + deployment: _Deployment, gateway: Gateway +) -> None: + with _listening(deployment) as peer: + response: Final = gateway.request( + "POST", + "/utils/token_counter", + {"model": deployment.model, "prompt": USER_TEXT}, + params={"call_endpoint": "true"}, + ) + payload: Final = _payload(response) + assert (payload["total_tokens"], payload["count_source"]) == (_COUNT, "provider"), response.text + assert payload["tokenizer_type"] == deployment.provider.tokenizer, response.text + assert _bodies(peer, deployment.provider) == ({"model": _MODEL, "messages": [USER]},) + + def test_responses_input_tokens_lifts_instructions(deployment: _Deployment, gateway: Gateway) -> None: with _listening(deployment) as peer: response: Final = gateway.request( diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index 2a0508402ae..07394c03f0c 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -550,6 +550,32 @@ def test_runtime_404_and_mantle_404_surface_provider_error(counting_proxy: Owned assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) +def test_utils_token_counter_prompt_answers_the_count_model_404_like_the_messages_form( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_rejecting(404)), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _metadata_deployment(scenario, runtime.url) + from_prompt: Final = gateway.request( + "POST", + "/utils/token_counter", + {"model": model, "prompt": "Count this message"}, + params={"call_endpoint": "true"}, + ) + from_messages: Final = gateway.request( + "POST", "/utils/token_counter", _bare(model), params={"call_endpoint": "true"} + ) + assert from_prompt.status_code == from_messages.status_code == 404, from_prompt.text + assert _REJECTION in from_prompt.text, from_prompt.text + assert _JSON_OBJECT.validate_json(from_prompt.content) == _JSON_OBJECT.validate_json(from_messages.content) + assert runtime.drain() == () + assert _mantle_bodies(mantle) == (_MANTLE_BARE, _MANTLE_BARE) + + def test_runtime_404_on_nova_does_not_call_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: gateway: Final = counting_proxy.gateway with ( diff --git a/tests/unit/proxy/router_entry_point_stubs.py b/tests/unit/proxy/router_entry_point_stubs.py index 306740489d7..10176accf9f 100644 --- a/tests/unit/proxy/router_entry_point_stubs.py +++ b/tests/unit/proxy/router_entry_point_stubs.py @@ -184,6 +184,7 @@ def route_entry_points_to_stubs(router: _RouterDouble) -> _RouterDouble: _core_get_settings=lambda: router.get_settings(), _core_get_available_guardrail=lambda name: router.get_available_guardrail(guardrail_name=name), _core_get_allowed_fails_from_policy=lambda exception: router.get_allowed_fails_from_policy(exception=exception), + _core_deployment_call_timeout=lambda litellm_params: router.deployment_call_timeout(litellm_params), ) router.settings.side_effect = lambda: RouterSettings( **{name: getattr(router, name) for name in SETTING_FIELDS}, _router=settings_backend diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 694b257b1f9..34868606012 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -8,10 +8,12 @@ import hashlib import itertools import json import logging +import socket import threading import traceback import uuid -from collections.abc import Callable, Coroutine, Mapping +from collections.abc import Callable, Coroutine, Iterator, Mapping +from contextlib import contextmanager from dataclasses import replace from functools import partial from pathlib import Path @@ -44,8 +46,11 @@ import litellm from litellm import Router from litellm._logging import JsonFormatter, verbose_proxy_logger from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler +from litellm.llms.azure_ai.anthropic.count_tokens import handler as azure_ai_count_tokens_http from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.bedrock.count_tokens import handler as bedrock_count_tokens_http +from litellm.llms.bedrock.count_tokens import mantle_handler as bedrock_mantle_count_tokens_http from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.kimi_k3 import ( @@ -4833,36 +4838,317 @@ async def test_local_count_of_a_prompt_with_system_or_tools_counts_them_like_the assert from_prompt.total_tokens == from_message.total_tokens > prompt_alone.total_tokens +_PROVIDER_COUNTED_DEPLOYMENTS: Final = ( + pytest.param( + {"model": "anthropic/claude-sonnet-5-5", "api_key": "test-key", "api_base": "https://count-tokens.test"}, + id="anthropic", + ), + pytest.param( + {"model": "azure_ai/claude-sonnet-5-5", "api_key": "test-key", "api_base": "https://count-tokens.test"}, + id="azure_ai", + ), + pytest.param( + {"model": "openai/gpt-5.6", "api_key": "test-key", "api_base": "https://count-tokens.test/v1"}, id="openai" + ), + pytest.param( + { + "model": "bedrock/us.anthropic.claude-sonnet-5-5", + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + }, + id="bedrock", + ), +) + + +@contextmanager +def _count_apis_answering(peer: Callable[[httpx.Request], httpx.Response]) -> Iterator[None]: + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + with ( + patch.object(anthropic_count_tokens_handler, "get_async_httpx_client", return_value=client), + patch.object(azure_ai_count_tokens_http, "get_async_httpx_client", return_value=client), + patch.object(bedrock_count_tokens_http, "get_async_httpx_client", return_value=client), + patch.object(bedrock_mantle_count_tokens_http, "get_async_httpx_client", return_value=client), + patch.object(openai_count_tokens_http, "get_async_httpx_client", return_value=client), + ): + yield + + +def _route_to_one_deployment(litellm_params: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "provider-counted", "litellm_params": litellm_params}]), + ) + + +_FOX_AS_PROMPT: Final = TokenCountRequest(model="provider-counted", prompt=_FOX) +_FOX_AS_MESSAGE: Final = TokenCountRequest(model="provider-counted", messages=[{"role": "user", "content": _FOX}]) + + @pytest.mark.asyncio -async def test_call_endpoint_prompt_keeps_the_local_count_on_a_counter_that_needs_messages( +@pytest.mark.parametrize("litellm_params", _PROVIDER_COUNTED_DEPLOYMENTS) +async def test_call_endpoint_sends_a_prompt_to_the_provider_as_one_user_message( + litellm_params: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + _route_to_one_deployment(litellm_params, monkeypatch) + sent_bodies: Final[list[JsonValue]] = [] # mutable-ok: collects what the mock peer receives + + def count_api(request: httpx.Request) -> httpx.Response: + sent_bodies.append(json.loads(request.content)) + return httpx.Response( + 200, json={"object": "response.input_tokens", "input_tokens": 41, "inputTokens": 41}, request=request + ) + + with _count_apis_answering(count_api): + from_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + from_message: Final = await token_counter(request=_FOX_AS_MESSAGE, call_endpoint=True) + + assert (from_prompt.total_tokens, from_prompt.error) == (41, False) + assert from_prompt.count_source != "local_estimate" + assert from_prompt == from_message + assert len(sent_bodies) == 2 and sent_bodies[0] == sent_bodies[1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [404, 503]) +@pytest.mark.parametrize("litellm_params", _PROVIDER_COUNTED_DEPLOYMENTS) +async def test_call_endpoint_answers_a_provider_error_on_a_prompt_like_on_the_same_text_sent_as_a_user_message( + litellm_params: dict[str, JsonValue], status: int, monkeypatch: pytest.MonkeyPatch +) -> None: + _route_to_one_deployment(litellm_params, monkeypatch) + + def count_api(request: httpx.Request) -> httpx.Response: + return httpx.Response(status, json={"error": {"message": f"scripted {status}"}}, request=request) + + with _count_apis_answering(count_api): + with pytest.raises(ProxyException) as from_prompt: + await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + with pytest.raises(ProxyException) as from_message: + await token_counter(request=_FOX_AS_MESSAGE, call_endpoint=True) + + assert from_prompt.value.code == str(status) + assert f"scripted {status}" in from_prompt.value.message + assert (from_prompt.value.code, from_prompt.value.type, from_prompt.value.message) == ( + from_message.value.code, + from_message.value.type, + from_message.value.message, + ) + + +@pytest.mark.asyncio +async def test_framed_stand_in_counts_a_prompt_like_the_same_text_sent_as_a_user_message( monkeypatch: pytest.MonkeyPatch, ) -> None: - router: Final = Router( - model_list=[ - { - "model_name": "anthropic-counted", - "litellm_params": { - "model": "anthropic/claude-sonnet-5-5", - "api_key": "test-key", - "api_base": "https://count-tokens.test", - }, - } - ] - ) - monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) - client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_unavailable_peer)) - request: Final = TokenCountRequest(model="anthropic-counted", prompt="The quick brown fox jumps over the lazy dog.") + _framed_kimi(monkeypatch) - with patch.object(anthropic_count_tokens_handler, "get_async_httpx_client", return_value=client): - with_endpoint: Final = await token_counter(request=request, call_endpoint=True) + from_prompt: Final = await token_counter(request=TokenCountRequest(model="kimi", prompt=_FOX), call_endpoint=True) + from_message: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": _FOX}]), call_endpoint=True + ) + local_prompt: Final = await token_counter(request=TokenCountRequest(model="kimi", prompt=_FOX)) + + assert (from_prompt.count_source, from_prompt.error) == ("local_estimate", False) + assert from_prompt.total_tokens == from_message.total_tokens > local_prompt.total_tokens + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", ["fireworks_ai/accounts/fireworks/models/glm-5", "deepseek/deepseek-v4-flash"], ids=["plain", "framed"] +) +async def test_call_endpoint_counts_a_prompt_on_a_provider_without_a_count_api_as_raw_text( + model: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _route_to_one_deployment({"model": model, "api_key": "test-key"}, monkeypatch) + sent: Final[list[httpx.Request]] = [] # mutable-ok: collects what the mock peer receives + + def count_api(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json={"input_tokens": 41}, request=request) + + with _count_apis_answering(count_api): + from_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + from_message: Final = await token_counter(request=_FOX_AS_MESSAGE, call_endpoint=True) + local_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=False) + + assert sent == [] + assert (from_prompt.count_source, from_prompt.error) == ("local_estimate", False) + assert from_prompt.total_tokens == local_prompt.total_tokens < from_message.total_tokens + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment", + [ + pytest.param( + {"litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "test-key"}}, id="cost map mode" + ), + pytest.param( + { + "litellm_params": {"model": "openai/in-house-embedder", "api_key": "test-key"}, + "model_info": {"mode": "embedding"}, + }, + id="configured mode", + ), + ], +) +async def test_call_endpoint_counts_a_prompt_on_a_non_chat_deployment_as_raw_text_without_the_provider( + deployment: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, "llm_router", Router(model_list=[{"model_name": "provider-counted", **deployment}]) + ) + sent: Final[list[httpx.Request]] = [] # mutable-ok: collects what the mock peer receives + + def count_api(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json={"input_tokens": 41}, request=request) + + with _count_apis_answering(count_api): + from_endpoint: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + local: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=False) + + assert sent == [] + assert (from_endpoint.total_tokens, from_endpoint.error) == (local.total_tokens, False) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["chat", "responses"]) +async def test_call_endpoint_sends_a_prompt_on_a_chat_mode_deployment_to_the_provider( + mode: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "provider-counted", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "test-key", + "api_base": "https://count-tokens.test/v1", + }, + "model_info": {"mode": mode}, + } + ] + ), + ) + sent: Final[list[httpx.Request]] = [] # mutable-ok: collects what the mock peer receives + + def count_api(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 41}, request=request) + + with _count_apis_answering(count_api): + from_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + + assert (from_prompt.total_tokens, from_prompt.error) == (41, False) + assert len(sent) == 1 + + +@pytest.mark.asyncio +async def test_call_endpoint_counts_a_prompt_for_a_model_without_a_deployment_as_raw_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "other", "litellm_params": {"model": "openai/gpt-5.6"}}]), + ) + request: Final = TokenCountRequest(model="deepseek/deepseek-v4-flash", prompt=_FOX) + + from_endpoint: Final = await token_counter(request=request, call_endpoint=True) local: Final = await token_counter(request=request, call_endpoint=False) - assert (with_endpoint.error, with_endpoint.total_tokens, with_endpoint.tokenizer_type) == ( - False, - local.total_tokens, - local.tokenizer_type, + assert (from_endpoint.total_tokens, from_endpoint.error) == (local.total_tokens, False) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "path"), + [("anthropic/claude-sonnet-5-5", ""), ("azure_ai/claude-sonnet-5-5", ""), ("openai/gpt-5.6", "/v1")], + ids=["anthropic", "azure_ai", "openai"], +) +@pytest.mark.parametrize( + ("row_timeout", "router_timeout"), + [({"timeout": 0.5}, None), ({"request_timeout": 0.5}, 30.0), ({}, 0.5)], + ids=["deployment timeout", "deployment request_timeout", "router_settings timeout"], +) +async def test_call_endpoint_answers_a_timeout_when_a_stalled_count_api_outlasts_the_chat_timeout( + model: str, + path: str, + row_timeout: dict[str, JsonValue], + router_timeout: float | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + with socket.create_server(("127.0.0.1", 0)) as never_answers: + port: Final = never_answers.getsockname()[1] + litellm_params: Final[dict[str, JsonValue]] = { + "model": model, + "api_key": "test-key", + "api_base": f"http://127.0.0.1:{port}{path}", + } + deployment: Final[dict[str, JsonValue]] = { + "model_name": "provider-counted", + "litellm_params": {**litellm_params, **row_timeout}, + } + monkeypatch.setattr( + litellm.proxy.proxy_server, "llm_router", Router(model_list=[deployment], timeout=router_timeout) + ) + + with pytest.raises(ProxyException) as error: + async with asyncio.timeout(10): + await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + + assert error.value.code == "408" + + +@pytest.mark.asyncio +async def test_responses_count_answers_a_timeout_when_a_stalled_count_api_outlasts_the_deployment_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with socket.create_server(("127.0.0.1", 0)) as never_answers: + port: Final = never_answers.getsockname()[1] + _route_to_one_deployment( + { + "model": "openai/gpt-5.6", + "api_key": "test-key", + "api_base": f"http://127.0.0.1:{port}/v1", + "timeout": 0.5, + }, + monkeypatch, + ) + + with pytest.raises(ProxyException) as error: + async with asyncio.timeout(10): + await count_request_tokens( + request=_FOX_AS_MESSAGE, + call_endpoint=True, + responses_request={"model": "provider-counted", "input": _FOX}, + ) + + assert error.value.code == "408" + + +@pytest.mark.asyncio +async def test_call_endpoint_names_the_status_of_a_count_api_error_with_an_empty_body( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _route_to_one_deployment( + {"model": "openai/gpt-5.6", "api_key": "test-key", "api_base": "https://count-tokens.test/v1"}, monkeypatch ) + with _count_apis_answering(lambda request: httpx.Response(404, request=request)): + with pytest.raises(ProxyException) as error: + await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + + assert error.value.code == "404" + assert "404" in error.value.message + _PINNED_COUNTS: Final = MappingProxyType({"deployment-a.test": 11, "deployment-b.test": 22})