mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(proxy): send a prompt-only call_endpoint token count to every provider count API (#50783)
* fix(proxy): count a prompt-only call_endpoint request with the provider on every counter POST /utils/token_counter?call_endpoint=true with only a prompt sent the prompt to the provider count API on Gemini and Vertex only. Anthropic, OpenAI, Bedrock and Azure AI (Claude) deployments answered a local estimate instead. Every counter now gets the prompt as one user message, so a provider error surfaces the same as the messages form, and the framed stand-in after a failed provider count frames the prompt as one user turn, matching the messages form * fix(proxy): count a call_endpoint prompt like a user message on providers without a count API * fix(proxy): count a call_endpoint prompt as a turn only on chat deployments A prompt-only /utils/token_counter request with call_endpoint=true now goes to the count API as one user message only when the deployment serves chat (mode chat, responses, or unknown). Embedding rows and models with no deployment keep the raw local count. The Anthropic, Azure AI, and OpenAI count calls now honor the deployment's timeout, and a count error with an empty body names its status code. * fix(proxy): fall back to router_settings.timeout for count calls and answer 408 on a count timeout Count calls to Anthropic, Azure AI, and OpenAI now pick their timeout the way a chat call to the same deployment does: the deployment's timeout, else router_settings.timeout, else request_timeout. A count call that times out answers 408 instead of a 500 processing error. * fix(proxy): count a prompt as a turn only where a count API exists and give count calls chat's timeout order * fix(proxy): read the count call timeout through the router settings snapshot The token counter now asks llm_router.settings().timeout_for() for the chat timeout order instead of a new Router member, so the router surface does not grow * fix(router): answer timeout_for on every router settings snapshot The proxy reads the count call timeout through llm_router.settings().timeout_for(), but only the Rust core's snapshot had that method, so a plain RouterSettings (the unit tests' router stub) raised AttributeError. timeout_for now lives on RouterSettings and reads the core's _core_deployment_call_timeout, like the snapshot's other members --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
paveln83
GitHub
parent
c1abd490b2
commit
90f37e113c
14 changed files with 427 additions and 36 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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()``)."""
|
||||
|
|
|
|||
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue