mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
fix(proxy): raise on provider count_tokens errors instead of falling back to tiktoken (#46121)
Co-authored-by: Krrish Dholakia <krrishdholakia@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yuneng-jiang <yuneng@berri.ai>
This commit is contained in:
Krrish Dholakia
GitHub
parent
784581bace
commit
719c44db3f
5 changed files with 110 additions and 113 deletions
|
|
@ -14116,19 +14116,12 @@ async def _try_provider_token_count(
|
|||
code=status_code,
|
||||
)
|
||||
if result is not None and result.error is True:
|
||||
if litellm.disable_token_counter is True:
|
||||
raise ProxyException(
|
||||
message=result.error_message or "Token counting failed",
|
||||
type="token_counting_error",
|
||||
param="model",
|
||||
code=result.status_code or 500,
|
||||
)
|
||||
verbose_proxy_logger.warning(
|
||||
"Provider token counting failed (%s): %s. Falling back to local tokenizer.",
|
||||
result.status_code,
|
||||
result.error_message,
|
||||
raise ProxyException(
|
||||
message=result.error_message or "Token counting failed",
|
||||
type="token_counting_error",
|
||||
param="model",
|
||||
code=result.status_code or 500,
|
||||
)
|
||||
return None
|
||||
return result
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -398,24 +398,24 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place(
|
|||
)
|
||||
|
||||
|
||||
def test_messages_count_tokens_falls_back_locally_when_every_message_is_system(
|
||||
def test_messages_count_tokens_surfaces_the_peer_error_when_every_message_is_system(
|
||||
deployment: _Deployment, gateway: Gateway
|
||||
) -> None:
|
||||
with _listening(deployment) as peer:
|
||||
body: Final[dict[str, JsonValue]] = {"model": deployment.model, "messages": [LEADING]}
|
||||
local: Final = _local_count(gateway, body)
|
||||
response: Final = _count(gateway, body)
|
||||
assert _payload(response) == {"input_tokens": local}, response.text
|
||||
assert response.status_code == 400, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
assert _bodies(peer, deployment.provider) == ({"model": _MODEL, "messages": [], "system": [LIFTED]},)
|
||||
|
||||
|
||||
def test_messages_count_tokens_leaves_the_request_untouched_for_a_non_text_system(
|
||||
def test_messages_count_tokens_surfaces_the_peer_error_for_a_non_text_system(
|
||||
deployment: _Deployment, gateway: Gateway
|
||||
) -> None:
|
||||
with _listening(deployment) as peer:
|
||||
local: Final = _local_count(gateway, count_request(deployment.model, STRING_CASE))
|
||||
response: Final = _count(gateway, {**count_request(deployment.model, STRING_CASE), "system": 5})
|
||||
assert _payload(response) == {"input_tokens": local}, response.text
|
||||
assert response.status_code == 400, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
assert _bodies(peer, deployment.provider) == ({"model": _MODEL, "messages": [LEADING, USER], "system": 5},)
|
||||
|
||||
|
||||
|
|
@ -428,14 +428,14 @@ def test_messages_count_tokens_repeated_request_lifts_each_time(deployment: _Dep
|
|||
assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * 2
|
||||
|
||||
|
||||
def test_messages_count_tokens_answers_a_leading_system_without_content_before_any_peer_call(
|
||||
def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_without_content(
|
||||
deployment: _Deployment, gateway: Gateway
|
||||
) -> None:
|
||||
with _listening(deployment) as peer:
|
||||
body: Final[dict[str, JsonValue]] = {"model": deployment.model, "messages": [{"role": "system"}, USER]}
|
||||
local: Final = _local_count(gateway, body)
|
||||
response: Final = _count(gateway, body)
|
||||
assert _payload(response) == {"input_tokens": local}, response.text
|
||||
assert 500 <= response.status_code < 600, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
assert peer.drain() == ()
|
||||
follow_up: Final = _count(gateway, count_request(deployment.model, STRING_CASE))
|
||||
assert _payload(follow_up) == {"input_tokens": _COUNT}, follow_up.text
|
||||
|
|
@ -458,14 +458,14 @@ def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value(
|
|||
|
||||
|
||||
@pytest.mark.parametrize("status", [400, 403, 404, 500, 503])
|
||||
def test_messages_count_tokens_falls_back_locally_when_the_peer_rejects_the_lifted_body(
|
||||
def test_messages_count_tokens_surfaces_the_peer_status_when_the_peer_rejects_the_lifted_body(
|
||||
deployment: _Deployment, gateway: Gateway, status: int
|
||||
) -> None:
|
||||
with _listening(deployment, _rejecting(status)) as peer:
|
||||
body: Final = count_request(deployment.model, STRING_CASE)
|
||||
local: Final = _local_count(gateway, body)
|
||||
response: Final = _count(gateway, body)
|
||||
assert _payload(response) == {"input_tokens": local}, response.text
|
||||
assert response.status_code == status, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
assert _bodies(peer, deployment.provider) == (_STRING_BODY,)
|
||||
|
||||
|
||||
|
|
@ -478,14 +478,13 @@ def test_messages_count_tokens_unauthenticated_request_never_reaches_the_peer(
|
|||
assert peer.drain() == ()
|
||||
|
||||
|
||||
def test_peer_outage_between_concurrent_waves_falls_back_then_recovers(
|
||||
def test_peer_outage_between_concurrent_waves_errors_then_recovers(
|
||||
deployment: _Deployment, gateway: Gateway
|
||||
) -> None:
|
||||
with ExitStack() as stack:
|
||||
clients: Final = _clients(stack, str(gateway.client.base_url), 8)
|
||||
pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients)))
|
||||
body: Final = count_request(deployment.model, STRING_CASE)
|
||||
local: Final = _local_count(gateway, body)
|
||||
|
||||
def count(client: httpx.Client) -> tuple[int, JsonValue]:
|
||||
return _counted_on(client, gateway.key, body)
|
||||
|
|
@ -493,7 +492,9 @@ def test_peer_outage_between_concurrent_waves_falls_back_then_recovers(
|
|||
with _listening(deployment) as peer:
|
||||
assert tuple(pool.map(count, clients)) == ((200, _COUNT),) * len(clients)
|
||||
assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * len(clients)
|
||||
assert tuple(pool.map(count, clients)) == ((200, local),) * len(clients)
|
||||
outage: Final = tuple(pool.map(count, clients))
|
||||
for status, tokens in outage:
|
||||
assert 500 <= status < 600 and tokens is None, outage
|
||||
with _listening(deployment) as revived:
|
||||
assert tuple(pool.map(count, clients)) == ((200, _COUNT),) * len(clients)
|
||||
assert _bodies(revived, deployment.provider) == (_STRING_BODY,) * len(clients)
|
||||
|
|
|
|||
|
|
@ -522,7 +522,7 @@ def test_model_the_runtime_counts_never_reaches_mantle(counting_proxy: OwnedProx
|
|||
assert mantle.drain() == ()
|
||||
|
||||
|
||||
def test_non_claude_model_the_runtime_cannot_count_falls_back_locally_without_mantle(
|
||||
def test_non_claude_model_the_runtime_cannot_count_surfaces_the_error_without_mantle(
|
||||
counting_proxy: OwnedProxy, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
|
|
@ -533,7 +533,8 @@ def test_non_claude_model_the_runtime_cannot_count_falls_back_locally_without_ma
|
|||
):
|
||||
model: Final = _deployment(scenario, runtime.url, _NOVA)
|
||||
response: Final = _count(gateway, _bare(model))
|
||||
assert _payload(response) == {"input_tokens": _local_count(gateway, _bare(model))}, response.text
|
||||
assert response.status_code == 400, response.text
|
||||
assert _UNSUPPORTED in response.text and "input_tokens" not in response.text, response.text
|
||||
assert _runtime_count_targets(runtime) == (f"/model/{_NOVA}/count-tokens",)
|
||||
assert mantle.drain() == ()
|
||||
|
||||
|
|
@ -551,17 +552,19 @@ def test_runtime_403_on_a_claude_model_never_reaches_mantle(counting_proxy: Owne
|
|||
"messages": [{"role": "user", "content": "status=403 Count this message"}],
|
||||
}
|
||||
response: Final = _count(gateway, body)
|
||||
assert _payload(response) == {"input_tokens": _local_count(gateway, body)}, response.text
|
||||
assert response.status_code == 403, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
assert len(_runtime_count_targets(runtime)) == 1
|
||||
assert mantle.drain() == ()
|
||||
|
||||
|
||||
def test_unreachable_runtime_falls_back_locally_without_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
|
||||
def test_unreachable_runtime_surfaces_the_error_without_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with wire_server(_mantle(), port=mantle_port) as mantle, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, _closed_port_url())
|
||||
response: Final = _count(gateway, _full(model))
|
||||
assert _payload(response) == {"input_tokens": _local_count(gateway, _full(model))}, response.text
|
||||
assert 500 <= response.status_code < 600, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
assert mantle.drain() == ()
|
||||
|
||||
|
||||
|
|
@ -770,7 +773,7 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place_for_mant
|
|||
)
|
||||
|
||||
|
||||
def test_messages_count_tokens_falls_back_locally_when_every_message_is_system_for_mantle(
|
||||
def test_messages_count_tokens_surfaces_the_mantle_error_when_every_message_is_system(
|
||||
counting_proxy: OwnedProxy, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
|
|
@ -781,14 +784,14 @@ def test_messages_count_tokens_falls_back_locally_when_every_message_is_system_f
|
|||
):
|
||||
model: Final = _deployment(scenario, runtime.url)
|
||||
body: Final[dict[str, JsonValue]] = {"model": model, "messages": [LEADING]}
|
||||
local: Final = _local_count(gateway, body)
|
||||
response: Final = _count(gateway, body)
|
||||
assert _payload(response) == {"input_tokens": local}, response.text
|
||||
assert response.status_code == 400, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
assert len(_runtime_count_targets(runtime)) == 1
|
||||
assert _mantle_bodies(mantle) == (_mantle_body(messages=[], system=[LIFTED]),)
|
||||
|
||||
|
||||
def test_messages_count_tokens_leaves_a_leading_system_in_place_beside_a_non_text_system_for_mantle(
|
||||
def test_messages_count_tokens_surfaces_the_mantle_error_beside_a_non_text_system(
|
||||
counting_proxy: OwnedProxy, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
|
|
@ -798,9 +801,9 @@ def test_messages_count_tokens_leaves_a_leading_system_in_place_beside_a_non_tex
|
|||
gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = _deployment(scenario, runtime.url)
|
||||
local: Final = _local_count(gateway, count_request(model, STRING_CASE))
|
||||
response: Final = _count(gateway, {**count_request(model, STRING_CASE), "system": 5})
|
||||
assert _payload(response) == {"input_tokens": local}, response.text
|
||||
assert response.status_code == 400, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
assert len(_runtime_count_targets(runtime)) == 1
|
||||
assert _mantle_bodies(mantle) == (_mantle_body(messages=[LEADING, USER], system=5),)
|
||||
|
||||
|
|
@ -821,7 +824,7 @@ def test_messages_count_tokens_repeated_request_lifts_the_leading_system_each_ti
|
|||
assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) * 2
|
||||
|
||||
|
||||
def test_messages_count_tokens_answers_a_leading_system_without_content_before_any_mantle_call(
|
||||
def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_without_content(
|
||||
counting_proxy: OwnedProxy, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
|
|
@ -832,9 +835,9 @@ def test_messages_count_tokens_answers_a_leading_system_without_content_before_a
|
|||
):
|
||||
model: Final = _deployment(scenario, runtime.url)
|
||||
body: Final[dict[str, JsonValue]] = {"model": model, "messages": [{"role": "system"}, USER]}
|
||||
local: Final = _local_count(gateway, body)
|
||||
response: Final = _count(gateway, body)
|
||||
assert _payload(response) == {"input_tokens": local}, response.text
|
||||
assert 500 <= response.status_code < 600, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
assert mantle.drain() == ()
|
||||
follow_up: Final = _count(gateway, count_request(model, STRING_CASE))
|
||||
assert _payload(follow_up) == {"input_tokens": _MANTLE_COUNT}, follow_up.text
|
||||
|
|
@ -889,7 +892,7 @@ def test_disabled_token_counter_counts_a_leading_system_through_mantle(
|
|||
assert len(_runtime_count_targets(runtime)) == 1
|
||||
|
||||
|
||||
def test_mantle_outage_between_concurrent_lifted_waves_falls_back_then_recovers(
|
||||
def test_mantle_outage_between_concurrent_lifted_waves_errors_then_recovers(
|
||||
counting_proxy: OwnedProxy, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
|
|
@ -900,7 +903,6 @@ def test_mantle_outage_between_concurrent_lifted_waves_falls_back_then_recovers(
|
|||
scenario: Final = stack.enter_context(gateway.scenario())
|
||||
model: Final = _deployment(scenario, runtime.url)
|
||||
body: Final = count_request(model, STRING_CASE)
|
||||
local: Final = _local_count(gateway, body)
|
||||
|
||||
def count(client: httpx.Client) -> tuple[int, JsonValue]:
|
||||
return _counted_on(client, gateway.key, body)
|
||||
|
|
@ -908,7 +910,9 @@ def test_mantle_outage_between_concurrent_lifted_waves_falls_back_then_recovers(
|
|||
with wire_server(_mantle(_strict), port=mantle_port) as mantle:
|
||||
assert tuple(pool.map(count, clients)) == ((200, _MANTLE_COUNT),) * len(clients)
|
||||
assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) * len(clients)
|
||||
assert tuple(pool.map(count, clients)) == ((200, local),) * len(clients)
|
||||
outage: Final = tuple(pool.map(count, clients))
|
||||
for status, tokens in outage:
|
||||
assert 500 <= status < 600 and tokens is None, outage
|
||||
with wire_server(_mantle(_strict), port=mantle_port) as revived:
|
||||
assert tuple(pool.map(count, clients)) == ((200, _MANTLE_COUNT),) * len(clients)
|
||||
assert _mantle_bodies(revived) == (_MANTLE_LIFTED,) * len(clients)
|
||||
|
|
@ -916,7 +920,7 @@ def test_mantle_outage_between_concurrent_lifted_waves_falls_back_then_recovers(
|
|||
|
||||
|
||||
@pytest.mark.parametrize("status", [400, 403, 404, 500, 503])
|
||||
def test_messages_count_tokens_falls_back_locally_when_mantle_errors(
|
||||
def test_messages_count_tokens_surfaces_the_mantle_status_when_mantle_errors(
|
||||
counting_proxy: OwnedProxy, mantle_port: int, status: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
|
|
@ -927,7 +931,8 @@ def test_messages_count_tokens_falls_back_locally_when_mantle_errors(
|
|||
):
|
||||
model: Final = _deployment(scenario, runtime.url)
|
||||
response: Final = _count(gateway, _full(model))
|
||||
assert _payload(response) == {"input_tokens": _local_count(gateway, _full(model))}, response.text
|
||||
assert response.status_code == status, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
assert len(_runtime_count_targets(runtime)) == 1
|
||||
assert _mantle_bodies(mantle) == (_MANTLE_FULL,)
|
||||
|
||||
|
|
@ -937,7 +942,7 @@ def test_messages_count_tokens_falls_back_locally_when_mantle_errors(
|
|||
[b'{"inputTokens": 7}', b"not json at all", b"{}"],
|
||||
ids=["runtime_key", "not_json", "empty_object"],
|
||||
)
|
||||
def test_messages_count_tokens_falls_back_locally_when_mantle_answers_without_input_tokens(
|
||||
def test_messages_count_tokens_errors_when_mantle_answers_without_input_tokens(
|
||||
counting_proxy: OwnedProxy, mantle_port: int, body: bytes
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
|
|
@ -948,21 +953,22 @@ def test_messages_count_tokens_falls_back_locally_when_mantle_answers_without_in
|
|||
):
|
||||
model: Final = _deployment(scenario, runtime.url)
|
||||
response: Final = _count(gateway, _bare(model))
|
||||
assert _payload(response) == {"input_tokens": _local_count(gateway, _bare(model))}, response.text
|
||||
assert 500 <= response.status_code < 600, response.text
|
||||
assert len(_runtime_count_targets(runtime)) == 1
|
||||
assert _mantle_bodies(mantle) == (_MANTLE_BARE,)
|
||||
|
||||
|
||||
def test_messages_count_tokens_falls_back_locally_when_mantle_is_unreachable(counting_proxy: OwnedProxy) -> None:
|
||||
def test_messages_count_tokens_errors_when_mantle_is_unreachable(counting_proxy: OwnedProxy) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with wire_server(_runtime) as runtime, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, runtime.url)
|
||||
response: Final = _count(gateway, _full(model))
|
||||
assert _payload(response) == {"input_tokens": _local_count(gateway, _full(model))}, response.text
|
||||
assert 500 <= response.status_code < 600, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
assert len(_runtime_count_targets(runtime)) == 1
|
||||
|
||||
|
||||
def test_messages_count_tokens_falls_back_locally_when_mantle_never_answers(
|
||||
def test_messages_count_tokens_errors_when_mantle_never_answers(
|
||||
counting_proxy: OwnedProxy, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
|
|
@ -983,18 +989,18 @@ def test_messages_count_tokens_falls_back_locally_when_mantle_never_answers(
|
|||
scenario: Final = stack.enter_context(gateway.scenario())
|
||||
model: Final = _deployment(scenario, runtime.url)
|
||||
body: Final = _full(model)
|
||||
local: Final = _local_count(gateway, body)
|
||||
future: Final = pool.submit(_counted_on, client, gateway.key, body)
|
||||
eventually(held.qsize, lambda size: size == 1, seconds=30)
|
||||
assert gateway.request("GET", "/health/liveliness").status_code == 200
|
||||
assert _local_count(gateway, body) == local
|
||||
assert _local_count(gateway, body) > 0
|
||||
assert not future.done()
|
||||
assert future.result(timeout=120) == (200, local)
|
||||
status, tokens = future.result(timeout=120)
|
||||
assert 500 <= status < 600 and tokens is None, (status, tokens)
|
||||
assert len(_runtime_count_targets(runtime)) == 1
|
||||
assert len(mantle.drain()) == 1
|
||||
|
||||
|
||||
def test_messages_count_tokens_falls_back_locally_when_mantle_rejects_a_non_text_system(
|
||||
def test_messages_count_tokens_surfaces_the_mantle_error_for_a_non_text_system(
|
||||
counting_proxy: OwnedProxy, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
|
|
@ -1005,7 +1011,8 @@ def test_messages_count_tokens_falls_back_locally_when_mantle_rejects_a_non_text
|
|||
):
|
||||
model: Final = _deployment(scenario, runtime.url)
|
||||
response: Final = _count(gateway, {**_bare(model), "system": 5})
|
||||
assert _payload(response) == {"input_tokens": _local_count(gateway, _bare(model))}, response.text
|
||||
assert response.status_code == 400, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
assert len(_runtime_count_targets(runtime)) == 1
|
||||
assert _mantle_bodies(mantle) == (_mantle_body(system=5),)
|
||||
|
||||
|
|
@ -1174,7 +1181,7 @@ def test_disabled_token_counter_surfaces_the_mantle_error_instead_of_counting_lo
|
|||
assert len(_runtime_count_targets(runtime)) == 4
|
||||
|
||||
|
||||
def test_mantle_outage_between_concurrent_waves_falls_back_then_recovers(
|
||||
def test_mantle_outage_between_concurrent_waves_errors_then_recovers(
|
||||
counting_proxy: OwnedProxy, mantle_port: int
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
|
|
@ -1185,7 +1192,6 @@ def test_mantle_outage_between_concurrent_waves_falls_back_then_recovers(
|
|||
scenario: Final = stack.enter_context(gateway.scenario())
|
||||
model: Final = _deployment(scenario, runtime.url)
|
||||
body: Final = _full(model)
|
||||
local: Final = _local_count(gateway, body)
|
||||
|
||||
def generate_then_count(client: httpx.Client) -> tuple[int, int, JsonValue]:
|
||||
return _generated_then_counted(client, gateway.key, model, body)
|
||||
|
|
@ -1197,7 +1203,8 @@ def test_mantle_outage_between_concurrent_waves_falls_back_then_recovers(
|
|||
assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _MANTLE_COUNT),) * len(clients)
|
||||
assert _mantle_bodies(mantle) == (_MANTLE_FULL,) * len(clients)
|
||||
outage: Final = tuple(pool.map(count_only, clients))
|
||||
assert outage == ((200, local),) * len(clients)
|
||||
for status, tokens in outage:
|
||||
assert 500 <= status < 600 and tokens is None, outage
|
||||
with wire_server(_mantle(), port=mantle_port) as revived:
|
||||
assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _MANTLE_COUNT),) * len(clients)
|
||||
assert _mantle_bodies(revived) == (_MANTLE_FULL,) * len(clients)
|
||||
|
|
|
|||
|
|
@ -344,15 +344,14 @@ def test_utils_token_counter_local_mode_never_calls_the_peer(gateway: Gateway) -
|
|||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def test_utils_token_counter_falls_back_locally_when_peer_rejects_openai_tools(gateway: Gateway) -> None:
|
||||
def test_utils_token_counter_surfaces_the_peer_error_for_openai_tools(gateway: Gateway) -> None:
|
||||
with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
body: Final = {**_bare(model), "tools": _OPENAI_TOOLS}
|
||||
response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "true"})
|
||||
payload: Final = _payload(response)
|
||||
assert response.status_code == 400, response.text
|
||||
assert "total_tokens" not in response.text, response.text
|
||||
assert _counted_bodies(wire) == ({**_PEER_BARE, "tools": _OPENAI_TOOLS},)
|
||||
assert payload["total_tokens"] == _local_count(gateway, body), response.text
|
||||
assert payload["tokenizer_type"] != "vertex_ai_partner_models", response.text
|
||||
|
||||
|
||||
def test_responses_input_tokens_counts_through_the_partner_peer(gateway: Gateway) -> None:
|
||||
|
|
@ -365,7 +364,7 @@ def test_responses_input_tokens_counts_through_the_partner_peer(gateway: Gateway
|
|||
assert _counted_bodies(wire) == (_PEER_BARE,)
|
||||
|
||||
|
||||
def test_responses_input_tokens_falls_back_locally_when_peer_rejects(gateway: Gateway) -> None:
|
||||
def test_responses_input_tokens_surfaces_the_peer_error(gateway: Gateway) -> None:
|
||||
with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
response: Final = gateway.request(
|
||||
|
|
@ -373,11 +372,10 @@ def test_responses_input_tokens_falls_back_locally_when_peer_rejects(gateway: Ga
|
|||
"/v1/responses/input_tokens",
|
||||
{"model": model, "input": "Count this message", "instructions": "Be terse", "tools": _RESPONSES_TOOLS},
|
||||
)
|
||||
payload: Final = _payload(response)
|
||||
assert response.status_code == 400, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
(sent,) = _counted_bodies(wire)
|
||||
assert sent["tools"] == _RESPONSES_TOOLS, sent
|
||||
local: Final = _local_count(gateway, {"model": model, "messages": sent["messages"], "tools": _RESPONSES_TOOLS})
|
||||
assert payload == {"object": "response.input_tokens", "input_tokens": local}, response.text
|
||||
|
||||
|
||||
def test_gemini_count_tokens_route_reaches_the_partner_peer(gateway: Gateway) -> None:
|
||||
|
|
@ -388,26 +386,26 @@ def test_gemini_count_tokens_route_reaches_the_partner_peer(gateway: Gateway) ->
|
|||
assert _counted_bodies(wire) == ({"model": _BACKEND, "messages": _GEMINI_MESSAGES},)
|
||||
|
||||
|
||||
def test_gemini_count_tokens_route_falls_back_locally_when_peer_rejects(gateway: Gateway) -> None:
|
||||
def test_gemini_count_tokens_route_surfaces_the_peer_error(gateway: Gateway) -> None:
|
||||
with wire_server(_peer(_rejecting(400))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY)
|
||||
payload: Final = _payload(response)
|
||||
assert response.status_code == 400, response.text
|
||||
assert "totalTokens" not in response.text, response.text
|
||||
assert _counted_bodies(wire) == ({"model": _BACKEND, "messages": _GEMINI_MESSAGES},)
|
||||
assert payload["totalTokens"] == _local_count(gateway, {"model": model, "messages": _GEMINI_MESSAGES})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", [400, 500, 503])
|
||||
def test_messages_count_tokens_falls_back_locally_when_peer_errors(gateway: Gateway, status: int) -> None:
|
||||
def test_messages_count_tokens_surfaces_the_peer_error_when_peer_errors(gateway: Gateway, status: int) -> None:
|
||||
with wire_server(_peer(_rejecting(status))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
response: Final = _count(gateway, _full(model))
|
||||
payload: Final = _payload(response)
|
||||
assert response.status_code == status, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
assert _counted_bodies(wire) == (_PEER_FULL,)
|
||||
assert payload == {"input_tokens": _local_count(gateway, _full(model))}, response.text
|
||||
|
||||
|
||||
def test_messages_count_tokens_falls_back_locally_when_token_endpoint_rejects(gateway: Gateway) -> None:
|
||||
def test_messages_count_tokens_surfaces_the_error_when_token_endpoint_rejects(gateway: Gateway) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.target == "/_oauth/token":
|
||||
return Reply(
|
||||
|
|
@ -421,17 +419,18 @@ def test_messages_count_tokens_falls_back_locally_when_token_endpoint_rejects(ga
|
|||
gateway, scenario, wire.url, vertex_credentials=service_account_json(_PROJECT, wire.url)
|
||||
)
|
||||
response: Final = _count(gateway, _full(model))
|
||||
payload: Final = _payload(response)
|
||||
assert 500 <= response.status_code < 600, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
targets: Final = frozenset(request.target for request in wire.drain())
|
||||
assert targets == {"/_oauth/token"}, targets
|
||||
assert payload == {"input_tokens": _local_count(gateway, _full(model))}, response.text
|
||||
|
||||
|
||||
def test_messages_count_tokens_falls_back_locally_when_peer_is_unreachable(gateway: Gateway) -> None:
|
||||
def test_messages_count_tokens_errors_when_peer_is_unreachable(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, _closed_port_url())
|
||||
response: Final = _count(gateway, _full(model))
|
||||
assert _payload(response) == {"input_tokens": _local_count(gateway, _full(model))}, response.text
|
||||
assert 500 <= response.status_code < 600, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -476,13 +475,13 @@ def test_messages_count_tokens_leaves_null_system_and_tools_out(gateway: Gateway
|
|||
assert _counted_bodies(wire) == (_PEER_BARE,)
|
||||
|
||||
|
||||
def test_messages_count_tokens_falls_back_locally_when_peer_rejects_a_non_text_system(gateway: Gateway) -> None:
|
||||
def test_messages_count_tokens_surfaces_the_peer_error_for_a_non_text_system(gateway: Gateway) -> None:
|
||||
with wire_server(_peer(_strict)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
response: Final = _count(gateway, {**_bare(model), "system": 5})
|
||||
payload: Final = _payload(response)
|
||||
assert response.status_code == 400, response.text
|
||||
assert "input_tokens" not in response.text, response.text
|
||||
assert _counted_bodies(wire) == ({**_PEER_BARE, "system": 5},)
|
||||
assert payload == {"input_tokens": _local_count(gateway, _bare(model))}, response.text
|
||||
|
||||
|
||||
def test_messages_count_tokens_forwards_a_5kb_system_verbatim(gateway: Gateway) -> None:
|
||||
|
|
@ -578,7 +577,7 @@ def test_disabled_token_counter_surfaces_provider_failures_instead_of_counting_l
|
|||
assert len(_counted_bodies(wire)) == 2
|
||||
|
||||
|
||||
def test_peer_outage_between_concurrent_waves_falls_back_then_recovers(gateway: Gateway) -> None:
|
||||
def test_peer_outage_between_concurrent_waves_errors_then_recovers(gateway: Gateway) -> None:
|
||||
with ExitStack() as stack:
|
||||
clients: Final = _clients(stack, str(gateway.client.base_url), 8)
|
||||
pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients)))
|
||||
|
|
@ -587,7 +586,6 @@ def test_peer_outage_between_concurrent_waves_falls_back_then_recovers(gateway:
|
|||
port: Final = int(wire.url.rsplit(":", 1)[1])
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
body: Final = _full(model)
|
||||
local: Final = _local_count(gateway, body)
|
||||
|
||||
def generate_then_count(client: httpx.Client) -> tuple[int, int, JsonValue]:
|
||||
return _generated_then_counted(client, gateway.key, model, body)
|
||||
|
|
@ -595,7 +593,8 @@ def test_peer_outage_between_concurrent_waves_falls_back_then_recovers(gateway:
|
|||
assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _PEER_COUNT),) * len(clients)
|
||||
assert _counted_bodies(wire) == (_PEER_FULL,) * len(clients)
|
||||
outage: Final = tuple(pool.map(lambda client: _counted_on(client, gateway.key, body), clients))
|
||||
assert outage == ((200, local),) * len(clients)
|
||||
for status, tokens in outage:
|
||||
assert 500 <= status < 600 and tokens is None, outage
|
||||
with wire_server(_peer(), port=port) as revived:
|
||||
assert tuple(pool.map(generate_then_count, clients)) == ((200, 200, _PEER_COUNT),) * len(clients)
|
||||
assert _counted_bodies(revived) == (_PEER_FULL,) * len(clients)
|
||||
|
|
|
|||
|
|
@ -310,20 +310,16 @@ async def test_internal_token_counter_anthropic_provider_detection():
|
|||
setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
|
||||
|
||||
# Test with is_direct_request=False (simulating call from Anthropic endpoint)
|
||||
response = await token_counter(
|
||||
request=TokenCountRequest(
|
||||
model="claude-test",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
),
|
||||
call_endpoint=True,
|
||||
)
|
||||
|
||||
print("Anthropic provider test response:", response)
|
||||
|
||||
# Verify response structure
|
||||
assert response.request_model == "claude-test"
|
||||
assert response.model_used == "claude-3-sonnet-20240229"
|
||||
assert response.total_tokens > 0
|
||||
# The provider counter is engaged: its API error surfaces instead of a local count
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await token_counter(
|
||||
request=TokenCountRequest(
|
||||
model="claude-test",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
),
|
||||
call_endpoint=True,
|
||||
)
|
||||
assert exc_info.value.type == "token_counting_error"
|
||||
|
||||
# Test with non-Anthropic provider
|
||||
llm_router = Router(
|
||||
|
|
@ -1018,10 +1014,10 @@ async def test_proxy_token_counter_error_raises_exception_when_disabled():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_token_counter_error_falls_back_when_enabled():
|
||||
async def test_proxy_token_counter_error_raises_when_enabled():
|
||||
"""
|
||||
Test that proxy token_counter falls back to local tokenizer when disable_token_counter=False
|
||||
and provider returns an error response.
|
||||
Test that proxy token_counter raises ProxyException when disable_token_counter=False
|
||||
and the provider returns an error response, never falling back to the local tokenizer.
|
||||
"""
|
||||
# Create error response
|
||||
error_response = TokenCountResponse(
|
||||
|
|
@ -1069,19 +1065,20 @@ async def test_proxy_token_counter_error_falls_back_when_enabled():
|
|||
mock_get_provider_token_counter
|
||||
)
|
||||
|
||||
# Should not raise, should fall back to local tokenizer
|
||||
result = await token_counter(
|
||||
request=TokenCountRequest(
|
||||
model="claude-bedrock",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
),
|
||||
call_endpoint=True,
|
||||
)
|
||||
with patch("litellm.token_counter") as mock_local_counter:
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await token_counter(
|
||||
request=TokenCountRequest(
|
||||
model="claude-bedrock",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
),
|
||||
call_endpoint=True,
|
||||
)
|
||||
|
||||
# Should have used the fallback tokenizer
|
||||
assert result.error is False
|
||||
assert result.total_tokens > 0
|
||||
assert result.tokenizer_type != "bedrock_api"
|
||||
assert exc_info.value.code == "429"
|
||||
assert "Rate limit exceeded" in exc_info.value.message
|
||||
assert exc_info.value.type == "token_counting_error"
|
||||
mock_local_counter.assert_not_called()
|
||||
finally:
|
||||
litellm.disable_token_counter = original_disable
|
||||
litellm.proxy.proxy_server._get_provider_token_counter = (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue